tidy up

Author: j-atkins

src/virtualship/instruments/adcp.py

@@ -11,6 +11,10 @@
     register_instrument,
 )
 
+# =====================================================
+# SECTION: Dataclass
+# =====================================================
+
 
 @dataclass
 class ADCP:
@@ -19,22 +23,34 @@ class ADCP:
     name: ClassVar[str] = "ADCP"
 
 
-# we specifically use ScipyParticle because we have many small calls to execute
-# there is some overhead with JITParticle and this ends up being significantly faster
+# =====================================================
+# SECTION: Particle Class
+# =====================================================
+
+
 _ADCPParticle = ScipyParticle.add_variables(
     [
         Variable("U", dtype=np.float32, initial=np.nan),
         Variable("V", dtype=np.float32, initial=np.nan),
     ]
 )
 
+# =====================================================
+# SECTION: Kernels
+# =====================================================
+
 
 def _sample_velocity(particle, fieldset, time):
     particle.U, particle.V = fieldset.UV.eval(
         time, particle.depth, particle.lat, particle.lon, applyConversion=False
     )
 
 
+# =====================================================
+# SECTION: InputDataset Class
+# =====================================================
+
+
 @register_input_dataset(InstrumentType.ADCP)
 class ADCPInputDataset(InputDataset):
     """Input dataset for ADCP instrument."""
@@ -70,6 +86,11 @@ def get_datasets_dict(self) -> dict:
         }
 
 
+# =====================================================
+# SECTION: Instrument Class
+# =====================================================
+
+
 @register_instrument(InstrumentType.ADCP)
 class ADCPInstrument(Instrument):
     """ADCP instrument class."""

src/virtualship/instruments/argo_float.py

@@ -17,6 +17,10 @@
 from virtualship.models.spacetime import Spacetime
 from virtualship.utils import register_input_dataset, register_instrument
 
+# =====================================================
+# SECTION: Dataclass
+# =====================================================
+
 
 @dataclass
 class ArgoFloat:
@@ -32,6 +36,10 @@ class ArgoFloat:
     drift_days: float
 
 
+# =====================================================
+# SECTION: Particle Class
+# =====================================================
+
 _ArgoParticle = JITParticle.add_variables(
     [
         Variable("cycle_phase", dtype=np.int32, initial=0.0),
@@ -48,6 +56,10 @@ class ArgoFloat:
     ]
 )
 
+# =====================================================
+# SECTION: Kernels
+# =====================================================
+
 
 def _argo_float_vertical_movement(particle, fieldset, time):
     if particle.cycle_phase == 0:
@@ -116,6 +128,11 @@ def _check_error(particle, fieldset, time):
         particle.delete()
 
 
+# =====================================================
+# SECTION: InputDataset Class
+# =====================================================
+
+
 @register_input_dataset(InstrumentType.ARGO_FLOAT)
 class ArgoFloatInputDataset(InputDataset):
     """Input dataset for ArgoFloat instrument."""
@@ -161,6 +178,11 @@ def get_datasets_dict(self) -> dict:
         }
 
 
+# =====================================================
+# SECTION: Instrument Class
+# =====================================================
+
+
 @register_instrument(InstrumentType.ARGO_FLOAT)
 class ArgoFloatInstrument(Instrument):
     """ArgoFloat instrument class."""

src/virtualship/instruments/base.py

@@ -246,12 +246,8 @@ def __init__(
 
     def load_input_data(self) -> FieldSet:
         """Load and return the input data as a FieldSet for the instrument."""
-        # TODO: can simulate_schedule.py be refactored to be contained in base.py and repsective instrument files too...?
         # TODO: tests need updating...!
 
-        #! TODO: E.g. ADCP is giving too much depth data?!
-        #! TODO: in fact output from most instruments doesn't look quite right...?
-
         try:
             data_dir = self._get_data_dir(self.directory)
             joined_filepaths = {
@@ -303,10 +299,11 @@ def run(self, measurements: list, out_path: str | Path) -> None:
                 spinner=ship_spinner,
             ) as spinner:
                 self.simulate(measurements, out_path)
-                spinner.ok("✅")
+                spinner.ok("✅\n")
         else:
             print(f"Simulating {self.name} measurements... ")
             self.simulate(measurements, out_path)
+            print("\n")
 
     def _get_data_dir(self, expedition_dir: Path) -> Path:
         space_time_region_hash = get_space_time_region_hash(

src/virtualship/instruments/ctd.py

@@ -13,6 +13,10 @@
 # TODO: add some kind of check that each instrument has a dataclass, particle class, InputDataset class and Instrument class?
 # TODO: probably as a test
 
+# =====================================================
+# SECTION: Dataclass
+# =====================================================
+
 
 @dataclass
 class CTD:
@@ -24,6 +28,10 @@ class CTD:
     max_depth: float
 
 
+# =====================================================
+# SECTION: Particle Class
+# =====================================================
+
 _CTDParticle = JITParticle.add_variables(
     [
         Variable("salinity", dtype=np.float32, initial=np.nan),
@@ -36,7 +44,9 @@ class CTD:
 )
 
 
-# TODO: way to group kernels together, just to make clearer?
+# =====================================================
+# SECTION: Kernels
+# =====================================================
 
 
 def _sample_temperature(particle, fieldset, time):
@@ -61,6 +71,11 @@ def _ctd_cast(particle, fieldset, time):
             particle.delete()
 
 
+# =====================================================
+# SECTION: InputDataset Class
+# =====================================================
+
+
 @register_input_dataset(InstrumentType.CTD)
 class CTDInputDataset(InputDataset):
     """Input dataset for CTD instrument."""
@@ -99,14 +114,17 @@ def get_datasets_dict(self) -> dict:
         }
 
 
+# =====================================================
+# SECTION: Instrument Class
+# =====================================================
+
+
 @register_instrument(InstrumentType.CTD)
 class CTDInstrument(Instrument):
     """CTD instrument class."""
 
     def __init__(self, expedition, directory):
         """Initialize CTDInstrument."""
-        #! TODO: actually don't need to download U and V for CTD simulation... can instead add mock/duplicate of T and name it U (also don't need V)!
-
         filenames = {
             "S": f"{CTD.name}_s.nc",
             "T": f"{CTD.name}_t.nc",

src/virtualship/instruments/ctd_bgc.py

@@ -10,6 +10,10 @@
 from virtualship.models.spacetime import Spacetime
 from virtualship.utils import add_dummy_UV, register_input_dataset, register_instrument
 
+# =====================================================
+# SECTION: Dataclass
+# =====================================================
+
 
 @dataclass
 class CTD_BGC:
@@ -21,6 +25,10 @@ class CTD_BGC:
     max_depth: float
 
 
+# =====================================================
+# SECTION: Particle Class
+# =====================================================
+
 _CTD_BGCParticle = JITParticle.add_variables(
     [
         Variable("o2", dtype=np.float32, initial=np.nan),
@@ -37,6 +45,10 @@ class CTD_BGC:
     ]
 )
 
+# =====================================================
+# SECTION: Kernels
+# =====================================================
+
 
 def _sample_o2(particle, fieldset, time):
     particle.o2 = fieldset.o2[time, particle.depth, particle.lat, particle.lon]
@@ -80,6 +92,11 @@ def _ctd_bgc_cast(particle, fieldset, time):
             particle.delete()
 
 
+# =====================================================
+# SECTION: InputDataset Class
+# =====================================================
+
+
 @register_input_dataset(InstrumentType.CTD_BGC)
 class CTD_BGCInputDataset(InputDataset):
     """Input dataset object for CTD_BGC instrument."""
@@ -143,6 +160,11 @@ def get_datasets_dict(self) -> dict:
         }
 
 
+# =====================================================
+# SECTION: Instrument Class
+# =====================================================
+
+
 @register_instrument(InstrumentType.CTD_BGC)
 class CTD_BGCInstrument(Instrument):
     """CTD_BGC instrument class."""

src/virtualship/instruments/drifter.py

@@ -10,6 +10,10 @@
 from virtualship.models.spacetime import Spacetime
 from virtualship.utils import register_input_dataset, register_instrument
 
+# =====================================================
+# SECTION: Dataclass
+# =====================================================
+
 
 @dataclass
 class Drifter:
@@ -21,6 +25,10 @@ class Drifter:
     lifetime: timedelta | None  # if none, lifetime is infinite
 
 
+# =====================================================
+# SECTION: Particle Class
+# =====================================================
+
 _DrifterParticle = JITParticle.add_variables(
     [
         Variable("temperature", dtype=np.float32, initial=np.nan),
@@ -30,6 +38,10 @@ class Drifter:
     ]
 )
 
+# =====================================================
+# SECTION: Kernels
+# =====================================================
+
 
 def _sample_temperature(particle, fieldset, time):
     particle.temperature = fieldset.T[time, particle.depth, particle.lat, particle.lon]
@@ -42,6 +54,11 @@ def _check_lifetime(particle, fieldset, time):
             particle.delete()
 
 
+# =====================================================
+# SECTION: InputDataset Class
+# =====================================================
+
+
 @register_input_dataset(InstrumentType.DRIFTER)
 class DrifterInputDataset(InputDataset):
     """Input dataset for Drifter instrument."""
@@ -82,6 +99,11 @@ def get_datasets_dict(self) -> dict:
         }
 
 
+# =====================================================
+# SECTION: Instrument Class
+# =====================================================
+
+
 @register_instrument(InstrumentType.DRIFTER)
 class DrifterInstrument(Instrument):
     """Drifter instrument class."""

src/virtualship/instruments/ship_underwater_st.py

@@ -8,6 +8,10 @@
 from virtualship.instruments.types import InstrumentType
 from virtualship.utils import add_dummy_UV, register_input_dataset, register_instrument
 
+# =====================================================
+# SECTION: Dataclass
+# =====================================================
+
 
 @dataclass
 class Underwater_ST:
@@ -16,13 +20,21 @@ class Underwater_ST:
     name: ClassVar[str] = "Underwater_ST"
 
 
+# =====================================================
+# SECTION: Particle Class
+# =====================================================
+
 _ShipSTParticle = ScipyParticle.add_variables(
     [
         Variable("S", dtype=np.float32, initial=np.nan),
         Variable("T", dtype=np.float32, initial=np.nan),
     ]
 )
 
+# =====================================================
+# SECTION: Kernels
+# =====================================================
+
 
 # define function sampling Salinity
 def _sample_salinity(particle, fieldset, time):
@@ -34,6 +46,11 @@ def _sample_temperature(particle, fieldset, time):
     particle.T = fieldset.T[time, particle.depth, particle.lat, particle.lon]
 
 
+# =====================================================
+# SECTION: InputDataset Class
+# =====================================================
+
+
 @register_input_dataset(InstrumentType.UNDERWATER_ST)
 class Underwater_STInputDataset(InputDataset):
     """Input dataset for Underwater_ST instrument."""
@@ -74,6 +91,11 @@ def get_datasets_dict(self) -> dict:
         }
 
 
+# =====================================================
+# SECTION: Instrument Class
+# =====================================================
+
+
 @register_instrument(InstrumentType.UNDERWATER_ST)
 class Underwater_STInstrument(Instrument):
     """Underwater_ST instrument class."""