roms-tools 3.0.0__py3-none-any.whl → 3.1.0__py3-none-any.whl

roms_tools/setup/cdr_release.py

@@ -3,7 +3,7 @@ from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from datetime import datetime
 from enum import StrEnum, auto
-from typing import Literal
+from typing import Annotated, Literal
 
 from annotated_types import Ge, Le
 from pydantic import (
@@ -15,7 +15,6 @@ from pydantic import (
     model_validator,
 )
 from pydantic_core.core_schema import ValidationInfo
-from typing_extensions import Annotated
 
 from roms_tools.setup.utils import get_tracer_defaults, get_tracer_metadata_dict
 
@@ -221,15 +220,23 @@ class Release(BaseModel):
     """
 
     name: str
+    """Unique identifier for the release."""
     lat: Annotated[float, Ge(-90), Le(90)]
+    """Latitude of the release location in degrees North."""
     lon: float
+    """Longitude of the release location in degrees East."""
     depth: NonNegativeFloat
+    """Depth of the release in meters."""
     hsc: NonNegativeFloat = 0.0
+    """Horizontal scale (standard deviation) of the release in meters."""
     vsc: NonNegativeFloat = 0.0
+    """Vertical scale (standard deviation) of the release in meters."""
     times: list[datetime]
+    """Time points of the release events."""
 
     # this should be defined by subclasses
     release_type: ReleaseType
+    """Type of the release."""
 
     model_config = ConfigDict(extra="forbid")
 
@@ -243,7 +250,6 @@ class Release(BaseModel):
         ValueError
             If times are not strictly increasing, or fall outside the [start_time, end_time] window.
         """
-
         if self.times and len(self.times) > 0:
             if not all(t1 < t2 for t1, t2 in zip(self.times, self.times[1:])):
                 raise ValueError(
@@ -257,7 +263,6 @@ class Release(BaseModel):
         Modifies `self.times` in place by prepending or appending times as needed.
         If `times` is empty, it will be set to [`start_time`, `end_time`].
         """
-
         if not self.times:
             self.times = [start_time, end_time]
         else:
@@ -275,7 +280,7 @@ class Release(BaseModel):
 class VolumeRelease(Release):
     """Represents a CDR release with volume flux and tracer concentrations.
 
-    Attributes
+    Parameters
     ----------
     name : str
         Unique identifier for the release.
@@ -332,19 +337,21 @@ class VolumeRelease(Release):
 
     times: list[datetime] = Field([])
     fill_values: Literal["auto", "zero"] = "auto"
+    """Strategy for filling missing tracer concentration values."""
     volume_fluxes: Flux | NonNegativeFloat | list[NonNegativeFloat] = Field(
         default=0.0, validate_default=True
     )
+    """Volume flux(es) of the release in m³/s over time."""
     tracer_concentrations: dict[
         str, Concentration | NonNegativeFloat | list[NonNegativeFloat]
     ] = Field({})
+    """Dictionary of tracer names and their concentration values."""
 
     release_type: Literal[ReleaseType.volume] = ReleaseType.volume
 
     @field_validator("tracer_concentrations", mode="after")
     @classmethod
     def _create_concentrations(cls, tracer_concentrations, info: ValidationInfo):
-
         defaults = get_tracer_defaults()
         for tracer_name in defaults.keys():
             if tracer_name in tracer_concentrations:
@@ -379,7 +386,6 @@ class VolumeRelease(Release):
 
     @model_validator(mode="after")
     def _check_lengths(self) -> "VolumeRelease":
-
         num_times = len(self.times)
 
         for tracer_concentrations in self.tracer_concentrations.values():
@@ -409,7 +415,6 @@ class VolumeRelease(Release):
     @model_serializer(mode="wrap")
     def _simplified_dump(self, pydantic_serializer) -> dict:
         """Return a simplified dict representation with flattened values."""
-
         # pydantic_serializer is a function that runs pydantic's default conversion
         # of a model to a dict. then, we make some custom modifications after that
         with warnings.catch_warnings():
@@ -433,7 +438,7 @@ class VolumeRelease(Release):
 class TracerPerturbation(Release):
     """Represents a CDR release with tracer fluxes and without any volume.
 
-    Attributes
+    Parameters
     ----------
     name : str
         Unique identifier for the release.
@@ -473,15 +478,15 @@ class TracerPerturbation(Release):
     tracer_fluxes: dict[str, Flux | NonNegativeFloat | list[NonNegativeFloat]] = Field(
         {}
     )
+    """Dictionary of tracer names and their non-negative flux values."""
 
-    release_type: Literal[
+    release_type: Literal[ReleaseType.tracer_perturbation] = (
         ReleaseType.tracer_perturbation
-    ] = ReleaseType.tracer_perturbation
+    )
 
     @field_validator("tracer_fluxes", mode="after")
     @classmethod
-    def create_fluxes(cls, tracer_fluxes):
-
+    def _create_fluxes(cls, tracer_fluxes):
         # Fill all tracer fluxes that are not provided with zero
         defaults = get_tracer_defaults()
         for tracer_name in defaults.keys():
@@ -495,7 +500,7 @@ class TracerPerturbation(Release):
         return tracer_fluxes
 
     @model_validator(mode="after")
-    def check_tracer_flux_lengths(self):
+    def _check_tracer_flux_lengths(self):
         num_times = len(self.times)
         for flux in self.tracer_fluxes.values():
             flux.check_length(num_times)
@@ -520,7 +525,6 @@ class TracerPerturbation(Release):
     @model_serializer(mode="wrap")
     def _simplified_dump(self, pydantic_serializer) -> dict:
         """Return a simplified dict representation with flattened values."""
-
         # pydantic_serializer is a function that runs pydantic's default conversion
         # of a model to a dict. then, we make some custom modifications after that
         with warnings.catch_warnings():