Fix satformer simulations in data extraction.

simonpf · Nov 28, 2024 · 4107677 · 4107677
1 parent f006f36
commit 4107677
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 10 deletions.
diff --git 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
@@ -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