rework of scene module started

mvcm/scene.py

@@ -4,11 +4,11 @@ import logging  # noqa
 
 import numpy as np
 import xarray as xr
+from attrs import Factory, define, field, validators
 
 import ancillary_data as anc
-
-# from typing import Dict
-
+from mvcm.constants import ConstantsNamespace as Constants
+from mvcm.constants import SceneConstants
 
 # lsf: land sea flag
 _scene_list = [
@@ -61,6 +61,22 @@ _flags = [
     "australia",
 ]
 
+land_water_flags = {
+    "shallow_ocean": 0,
+    "land": 1,
+    "coastline": 2,
+    "shallow_inland": 3,
+    "ephemeral": 4,
+    "deep_inland": 5,
+    "continental": 6,
+    "deep_ocean": 7,
+    "no_data": 255,
+}
+
+ecosystem_flags = {
+    "water": 14,
+}
+
 # temp value, need to verify what the actual bad_data value is in the C code
 _bad_data = -999.0
 
@@ -123,6 +139,131 @@ def test_scene():
 """
 
 
+def compute_ref_angle(vza, sza, raz):
+    """Compute ref angle."""
+    cos_refang = np.sin(vza * Constants.DTR) * np.sin(sza * Constants.DTR) * np.cos(
+        raz * Constants.DTR
+    ) + np.cos(vza * Constants.DTR) * np.cos(sza * Constants.DTR)
+    return np.arccos(cos_refang) * Constants.RTD
+
+
+@define(kw_only=True, slots=True)
+class Scene:
+    """Class that defines scene type based on pixel information."""
+
+    data: xr.Dataset = field(validator=[validators.instance_of(xr.Dataset)])
+    # sunglint_angle: float
+    # snow_mask_thresholds: dict | None = None
+
+    ref_ang: np.ndarray = field(
+        init=False,
+        default=Factory(
+            lambda self: compute_ref_angle(
+                self.data.sensor_zenith.values,
+                self.data.solar_zenith.values,
+                self.data.relative_azimuth.values,
+            ),
+            takes_self=True,
+        ),
+    )
+
+    out_shape: tuple = field(
+        init=False, default=Factory(lambda self: self.latitude.shape, takes_self=True)
+    )
+
+    def find_day(self):
+        """Find daytime pixels."""
+        return xr.where(self.data.solar_zenith <= SceneConstants.SOLAR_ZENITH_DAY, 1, 0)
+
+    def find_visusd(self):
+        """Find where daytime pixels are used."""
+        return xr.where(self.data.solar_zenith <= SceneConstants.SOLAR_ZENITH_DAY, 1, 0)
+
+    def find_night(self):
+        """Find nightime pixels."""
+        return xr.where(self.data.solar_zenith > SceneConstants.SOLAR_ZENITH_DAY, 1, 0)
+
+    def find_polar(self):
+        """Find polar pixels."""
+        return xr.where(np.abs(self.data.latitude) > SceneConstants.POLAR_LATITUDES, 1, 0)
+
+    def find_coast(self):
+        """Find coast pixels."""
+        pass
+
+    def find_shallow_lake(self):
+        """Find shallow lake pixels."""
+        pass
+
+    def find_shallow_ocean(self):
+        """Find shallow ocean pixels."""
+        pass
+
+    def find_water(self):
+        """Find water pixels."""
+        pass
+
+    def find_snow(self):
+        """Find snow pixels."""
+        pass
+
+    def find_ice(self):
+        """Find ice pixels."""
+        pass
+
+    def find_new_zealand(self):
+        """Find new zealand pixels."""
+        pass
+
+    def find_uniform(self):
+        """Find uniform pixels."""
+        pass
+
+    def find_australia(self):
+        """Find australia pixels."""
+        pass
+
+    def find_desert(self):
+        """Find desert pixels."""
+        pass
+
+    def find_land(self):
+        """Find land pixels."""
+        pass
+
+    def find_sunglint(self):
+        """Find sun glint pixels."""
+        pass
+
+    def find_greenland(self):
+        """Find greenland pixels."""
+        pass
+
+    def find_high_elevation(self):
+        """Find high elevation pixels."""
+        pass
+
+    def find_antarctica(self):
+        """Find antarctica pixels."""
+        pass
+
+    def find_vrused(self):
+        """Find vrused pixels."""
+        pass
+
+    def find_map_snow(self):
+        """Find map snow pixels."""
+        pass
+
+    def find_map_ice(self):
+        """Find map ice pixels."""
+        pass
+
+    def find_ndsi_snow(self):
+        """Find ndsi snow pixels."""
+        pass
+
+
 def find_scene(
     data: xr.Dataset, sunglint_angle: float, snow_mask_thresholds: dict | None = None
 ) -> dict:
@@ -150,6 +291,8 @@ def find_scene(
 
     # Define constants to use in the scene classification
     solar_zenith_day = 85
+    polar_latitudes = 60
+    polar_coast_latitudes = 64
 
     day = np.zeros((dim1, dim2))
     day[sza <= solar_zenith_day] = 1
@@ -158,55 +301,89 @@ def find_scene(
     ) * np.cos(sza * _dtr)
     refang = np.arccos(cos_refang) * _rtd
 
-    # tmp = np.ones((dim1, dim2))
-    # tmp[day == 1] = day
-    # tmp[day == 0] = night
-
     scene_flag = {flg: np.zeros((dim1, dim2), dtype=np.int8) for flg in _flags}
+    # scene_flag = xr.Dataset()
+
+    # scene_flag = xr.Dataset(
+    #     {
+    #         flg: xr.full_like(data.latitude, fill_value=0, dtype=np.int8, chunks="auto")
+    #         for flg in _flags
+    #     }
+    # )
+
+    # def find_day(data: xr.Dataset, solar_zenith_day: np.float = 40) -> xr.DataArray:
+    #     """Find daytime pixels."""
+    #     return xr.where(data.solar_zenith_day < solar_zenith_day, 1, 0)
+
+    scene_flag["day"] = xr.where(data.solar_zenith <= solar_zenith_day, 1, 0)
+    scene_flag["visusd"] = xr.where(data.solar_zenith <= solar_zenith_day, 1, 0)
+    scene_flag["night"] = xr.where(data.solar_zenith > solar_zenith_day, 1, 0)
+
+    scene_flag["polar"] = xr.where(np.abs(data.latitude) > polar_latitudes, 1, 0)
 
-    tmp = xr.full_like(data.latitude, fill_value=1)
-    tmp = tmp.where(sza <= solar_zenith_day, other=0)
-    scene_flag["day"] = tmp.values
-    tmp = xr.full_like(data.latitude, fill_value=1)
-    tmp = tmp.where(sza <= solar_zenith_day, other=0)
-    scene_flag["visusd"] = tmp.values
-    tmp = xr.full_like(data.latitude, fill_value=1)
-    tmp = tmp.where(sza > solar_zenith_day, other=0)
-    scene_flag["night"] = tmp.values
-
-    # scene_flag["day"][sza <= solar_zenith_day] = 1
-    # scene_flag["visusd"][sza <= solar_zenith_day] = 1
-    # scene_flag["night"][sza > solar_zenith_day] = 1
-
-    scene_flag["polar"][np.abs(lat) > 60] = 1
     # ################# need to pass refang (once I figure out what it is)
     # and sunglint_angle. The latter
     # comes from the thresholds file. In the C code is snglnt_bounds[3]
-    idx = np.nonzero((scene_flag["day"] == 1) & (refang <= sunglint_angle))
-    scene_flag["sunglint"][idx] = 1
+    scene_flag["sunglint"] = xr.where((scene_flag["day"] == 1) & (refang <= sunglint_angle), 1, 0)
 
     # Force consistency between lsf and ecosystem type for water
-    idx = np.nonzero((lsf == 0) | ((lsf >= 5) & (lsf < 7)))
-    eco[idx] = 14
+    # idx = np.nonzero((lsf == 0) | ((lsf >= 5) & (lsf < 7)))
+    # eco[idx] = 14
+    data.ecosystem.values = data.ecosystem.where(
+        ~(
+            (data.land_water_mask == land_water_flags["shallow_ocean"])
+            | (data.land_water_mask == land_water_flags["deep_inland"])
+            | (data.land_water_mask == land_water_flags["deep_ocean"])
+        ),
+        14,
+    )
 
     # start by defining everything as land
-    scene_flag["land"] = np.ones((dim1, dim2), dtype=np.int8)
+    # scene_flag["land"] = np.ones((dim1, dim2), dtype=np.int8)
     scene_flag["water"] = np.zeros((dim1, dim2), dtype=np.int8)
+    scene_flag["land"] = xr.ones_like(data.latitude, dtype=np.int8)
+    # scene_flag["water"] = xr.zeros_like(data.latitude, dtype=np.int8)
 
     # Fix-up for missing ecosystem data in eastern Greenland and north-eastern Siberia.
     # Without this, these regions become completely "coast".
-    idx = np.nonzero((lsf != 255) & ((lsf == 1) | (lsf == 4)))
-    scene_flag["land"][idx] = 1
-
-    # idx = np.nonzero((lsf != 255) & (eco == 14))
-
-    idx = np.nonzero(((lsf == 1) | (lsf == 4)) & (eco == 14) & (lat < 64.0))
-    scene_flag["coast"][idx] = 1
+    # idx = np.nonzero((lsf != 255) & ((lsf == 1) | (lsf == 4)))
+    # scene_flag["land"][idx] = 1
+    scene_flag["land"] = xr.where(
+        (data.land_water_mask != land_water_flags["no_data"])
+        & (data.land_water_mask != land_water_flags["land"])
+        & (data.land_water_mask != land_water_flags["ephemeral"]),
+        1,
+        scene_flag["land"],
+    ).values
+
+    # idx = np.nonzero(((lsf == 1) | (lsf == 4)) & (eco == 14) & (lat < 64.0))
+    # scene_flag["coast"][idx] = 1
+    scene_flag["coast"] = xr.where(
+        (
+            (data.land_water_mask == land_water_flags["land"])
+            | (data.land_water_mask == land_water_flags["ephemeral"])
+        )
+        & (data.ecosystem == ecosystem_flags["water"])
+        & (data.latitude < polar_coast_latitudes),
+        1,
+        0,
+    )
+
+    # idx = np.nonzero(
+    #     (eco == 14) & ((lsf == 1) | (lsf == 4)) & ((lat >= 67.5) & (lon < -40.0) & (lon > -168.6))
+    # )
+    # scene_flag["coast"][idx] = 1
+    scene_flag["coast"] = xr.where(
+        (
+            (data.ecosystem == ecosystem_flags["water"])
+            & (data.land_water_mask == land_water_flags["land"])
+            | (data.land_water_mask == land_water_flags["ephemeral"])
+        )
+        & ((data.latitude >= 67.5) & (data.longitude < -40.0) & (data.longitude > -168.6)),
+        1,
+        scene_flag["coast"],
+    ).values
 
-    idx = np.nonzero(
-        (eco == 14) & ((lsf == 1) | (lsf == 4)) & ((lat >= 67.5) & (lon < -40.0) & (lon > -168.6))
-    )
-    scene_flag["coast"][idx] = 1
     idx = np.nonzero((eco == 14) & ((lsf == 1) | (lsf == 4)) & ((lat >= 67.5) & (lon > -12.5)))
     scene_flag["coast"][idx] = 1