Merge pull request #59 from axiom-data-science/schema-improvements

JSON Schema oil type improvements and related updates and refactors

Author: honzaflash

.pre-commit-config.yaml

@@ -54,4 +54,4 @@ repos:
     - id: codespell
       args:
         - --quiet-level=2
-      exclude: 'particle_tracking_manager/models/opendrift/enums.py'
+      exclude: 'particle_tracking_manager/models/opendrift/enums/oil_types.py'

docs/environment.yml

@@ -10,11 +10,12 @@ dependencies:
    - appdirs
    - cmocean
    - fastparquet
+   - importlib_resources  # should update past opendrift 1.12 and remove this when possible
    - numpy
    - xarray
    - kerchunk==0.2.7
    - numpydoc
-   - opendrift==1.13.0
+   - opendrift <=1.13.0
    # These are needed for the docs themselves
    - furo
    - jupytext

docs/tutorial.md

@@ -211,7 +211,7 @@ There are also specific seeding options for this model:
 
 ```{code-cell} ipython3
 m = ptm.OpenDriftModel(drift_model="OpenOil", lon=-89.85, lat=28.08, number=10, steps=45,
-                       z=-10, do3D=True, oil_type=('EC00561', 'Cook Inlet [2003]'),
+                       z=-10, do3D=True, oil_type="EC00561",
                        start_time="2009-11-19T12:00", ocean_model="TXLA",
                        ocean_model_local=False,
                        )
@@ -220,13 +220,24 @@ m.o.set_config('environment:constant:x_wind', -1)
 m.o.set_config('environment:constant:y_wind', 1)
 ```
 
-List available oil types from NOAA's ADIOS database (or check the database online):
+Note that `oil_type` was input by its oil id, in order to disambiguate the oils in ADIOS.
+
+List of available oil types can be found in the
+1. NOAA's ADIOS database https://adios.orr.noaa.gov/oils
+2. or in the library's `OIL_ID_TO_NAME` dictionary:
+    ```{code-cell} ipython3
+    ptm.models.opendrift.enums.oil_types.OIL_ID_TO_NAME
+    ```
+
+You can also find the oil IDs by name in said [database](https://adios.orr.noaa.gov/oils)
+or in the library's `NAME_TO_OIL_ID` dictionary. E.g.:
 
 ```{code-cell} ipython3
-ptm.OpenOilModelConfig.model_json_schema()["$defs"]["OilTypeEnum"]["enum"]
+ptm.models.opendrift.enums.oil_types.NAME_TO_OIL_ID["ALASKA NORTH SLOPE"]
 ```
 
-Note that `oil_type` was input as a tuple containing the oil id and name, in order to disambiguate the oils in ADIOS. You can instead refer to the `oil_type` by its id, for example in this case `oil_tye='EC00561'` would work.
+Keep in mind that some oil types share the same name which is why `NAME_TO_OIL_ID`
+contains lists of IDs.
 
 The configuration parameters for this simulation are:
 

docs/whats_new.md

@@ -2,8 +2,12 @@
 
 ## Unreleased
 
-* Fix kerchunk JSON file filtering.
+* Fixed kerchunk JSON file filtering.
 * Added CIOFS3 as ocean model option.
+* Refactored `OilTypeEnum` so that it's value is only the oil ID. Also add the
+  "oneOf" json schema mapping directly to the field's `json_schema_extra` so that
+  it is part of the schema by default.
+
 
 ## v0.12.2 (April 15, 2025)
 

particle_tracking_manager/__init__.py

@@ -10,12 +10,24 @@
 logging.getLogger("numcodecs").setLevel(logging.WARNING)
 
 from .config_the_manager import TheManagerConfig
+from .models import opendrift as opendrift_models
 from .models.opendrift.config_opendrift import (
     LarvalFishModelConfig,
     LeewayModelConfig,
     OceanDriftModelConfig,
     OpenDriftConfig,
     OpenOilModelConfig,
 )
-from .models.opendrift.enums import ModifyOilTypeJsonSchema
 from .models.opendrift.opendrift import OpenDriftModel
+
+
+__all__ = [
+    "TheManagerConfig",
+    "LarvalFishModelConfig",
+    "LeewayModelConfig",
+    "OceanDriftModelConfig",
+    "OpenDriftConfig",
+    "OpenOilModelConfig",
+    "OpenDriftModel",
+    "opendrift_models",
+]

particle_tracking_manager/config_ocean_model.py

@@ -129,7 +129,7 @@ def open_dataset(self, drop_vars: list) -> xr.Dataset:
 
 def create_ocean_model_simulation(
     ocean_model: OceanModelConfig,
-) -> OceanModelSimulation:
+) -> type[OceanModelSimulation]:
     """Create an ocean model simulation object."""
     ocean_model_name = ocean_model.name
     simulation_model = create_model(

particle_tracking_manager/config_the_manager.py

@@ -232,7 +232,8 @@ class TheManagerConfig(BaseModel):
     )
 
     model_config = {
-        "validate_defaults": True,
+        "validate_default": True,
+        # Field values will be returned as the `enum.value` - a string in most our cases
         "use_enum_values": True,
         "extra": "forbid",
     }

particle_tracking_manager/models/__init__.py

@@ -1 +1,3 @@
 """Options for models."""
+
+from .opendrift import *

particle_tracking_manager/models/opendrift/config_opendrift.py

@@ -7,7 +7,7 @@
 from pathlib import Path
 
 # Third-party imports
-from pydantic import Field, field_validator, model_validator
+from pydantic import Field, model_validator
 from pydantic.fields import FieldInfo
 from typing_extensions import Self
 
@@ -18,7 +18,6 @@
     DiffusivityModelEnum,
     DriftModelEnum,
     DropletSizeDistributionEnum,
-    ModifyOilTypeJsonSchema,
     ObjectTypeEnum,
     OilTypeEnum,
     PlotTypeEnum,
@@ -158,13 +157,6 @@ class OpenDriftConfig(TheManagerConfig):
         "extra": "forbid",
     }
 
-    @classmethod
-    def model_json_schema(cls, **kwargs) -> dict:
-        """Override the method to customize the JSON schema to include customization of oil_type."""
-        return super().model_json_schema(
-            schema_generator=ModifyOilTypeJsonSchema, **kwargs
-        )
-
     @model_validator(mode="after")
     def check_interpolator_filename(self) -> Self:
         """Check if interpolator_filename is set correctly."""
@@ -489,10 +481,14 @@ class OpenOilModelConfig(OceanDriftModelConfig):
     drift_model: DriftModelEnum = DriftModelEnum.OpenOil  # .value
 
     oil_type: OilTypeEnum = Field(
-        default=OilTypeEnum["AD04012"],  # .value,
+        default=OilTypeEnum.AD04012,  # .value,
         description="Oil type to be used for the simulation, from the NOAA ADIOS database.",
         title="Oil Type",
-        json_schema_extra={"od_mapping": "seed:oil_type", "ptm_level": 1},
+        json_schema_extra={
+            "od_mapping": "seed:oil_type",
+            "ptm_level": 1,
+            "oneOf": [{"const": oil.value, "title": oil.title} for oil in OilTypeEnum],
+        },
     )
 
     m3_per_hour: float = Field(
@@ -645,20 +641,6 @@ class OpenOilModelConfig(OceanDriftModelConfig):
         OceanDriftModelConfig.model_fields["vertical_mixing"], Field(default=True)
     )
 
-    @field_validator("oil_type", mode="before")
-    def map_id_to_oil_type_tuple(cls, v):
-        """Map input oil type to enum name (which is the oil type id)."""
-        if (
-            v in OilTypeEnum.__members__
-        ):  # Check if it matches an Enum name (which is the oil type id)
-            return OilTypeEnum[v]  # then return id
-        else:
-            return v  # return the original value
-        # for enum_member in OilTypeEnum:  # Check if it matches an Enum value (which is the oil type name/title)
-        #     if enum_member.value == v:
-        #         return enum_member  # then return id
-        raise ValueError(f"Invalid value or name '{v}' for OilTypeEnum")
-
 
 class LarvalFishModelConfig(OceanDriftModelConfig):
     """Larval fish model configuration for OpenDrift."""
@@ -715,8 +697,8 @@ class LarvalFishModelConfig(OceanDriftModelConfig):
         },
     )
 
-    length: float = Field(
-        default=0,
+    length: float | None = Field(
+        default=None,
         description="Seeding value of length. This is not currently used.",
         title="Length",
         gt=0,