diff --git a/gprof_nn/data/sim.py b/gprof_nn/data/sim.py index 0977b9c..caf15fa 100644 --- a/gprof_nn/data/sim.py +++ b/gprof_nn/data/sim.py @@ -20,6 +20,7 @@ from typing import Optional, Tuple import click +from filelock import FileLock import numpy as np import pandas as pd from pansat import TimeRange, Granule @@ -66,6 +67,7 @@ calculate_obs_properties, RADIUS_OF_INFLUENCE, calculate_polarization_weights, + decompress_scene ) BEAM_WIDTHS = { @@ -645,7 +647,7 @@ def __iter__(self): if self.target_sensor.kind in ["CONICAL", "CONICAL_CONSTELLATION"]: yield self.load_input_data_conical() elif self.target_sensor.kind in ["XTRACK"]: - yield self.load_input_data_conical() + yield self.load_input_data_xtrack() else: raise ValueError(f"Encountered sensor with unsupported kind {self.target_sensor.kind}.") @@ -937,14 +939,16 @@ def process_sim_file( data = decompress_scene(data, vars) if satformer_model is not None: - simulate_tbs_satformer(satformer_model, sim_data, sensors.AMSR2) + simulate_tbs_satformer(satformer_model, data, sensor) if sensor.name != "GMI": - data = simulate_tbs_satformer( - model_path, - data, - sensor, - ) + lock = FileLock("cuda.lock") + with lock: + simulate_tbs_satformer( + satformer_model, + data, + sensor, + ) if lonlat_bounds is not None: lon_ll, lat_ll, lon_ur, lat_ur = lonlat_bounds diff --git a/gprof_nn/data/utils.py b/gprof_nn/data/utils.py index 5d5f70d..885b461 100644 --- a/gprof_nn/data/utils.py +++ b/gprof_nn/data/utils.py @@ -290,6 +290,13 @@ def save_scene( "_FillValue": 2 ** 16 - 1, "zlib": True }, + "satformer_tbs": { + "dtype": "uint16", + "scale_factor": 0.01, + "add_offset": 1, + "_FillValue": 2 ** 16 - 1, + "zlib": True + }, "brightness_temperature_biases": { "dtype": "int16", "scale_factor": 0.01, @@ -828,7 +835,7 @@ def calculate_angles( BEAM_WIDTHS = { "gmi": [1.75, 1.75, 1.0, 1.0, 0.9, 0.9, 0.9, 0.4, 0.4, 0.4, 0.4, 0.4, 0.4], - "atms": [5.2, 5.2, 2.2, 1.1, 1.1, 1.1, 1.1, 1.1], + "atms": [5.2, 5.2, 2.2, 1.1, 1.1, 1.1, 1.1, 1.1, 1.1], "amsr2": [1.2, 1.2, 0.65, 0.65, 0.75, 0.75, 0.35, 0.35, 0.15, 0.15, 0.15, 0.15], } @@ -874,10 +881,11 @@ def calculate_obs_properties( l1c_file = L1CFile(granule.file_record.local_path) sensor = l1c_file.sensor.name.lower() + tot_chan_ind = 0 + granule_data = granule.open() if "latitude" in granule_data: pass - else: swath_ind = 1 while f"latitude_s{swath_ind}" in granule_data: @@ -945,13 +953,14 @@ def calculate_obs_properties( freqs[chan_ind] * np.ones_like(observations[-1]), offsets[chan_ind] * np.ones_like(observations[-1]), calculate_polarization_weights(pols[chan_ind], viewing_angle), - BEAM_WIDTHS[sensor][chan_ind] * np.ones_like(observations[-1]), + BEAM_WIDTHS[sensor][tot_chan_ind] * np.ones_like(observations[-1]), sensor_alt, zenith, 1.0 + np.sin(np.deg2rad(azimuth)), 1.0 + np.cos(np.deg2rad(azimuth)) )) meta_data.append(meta) + tot_chan_ind += 1 swath_ind += 1