finish gravity uijson and round trip tests to file and back

Author: benk-mira

simpeg_drivers-assets/uijson/gravity_inversion.ui.json

@@ -19,9 +19,7 @@
             "{6a057fdc-b355-11e3-95be-fd84a7ffcb88}",
             "{f26feba3-aded-494b-b9e9-b2bbcbe298e1}",
             "{48f5054a-1c5c-4ca4-9048-80f36dc60a06}",
-            "{b020a277-90e2-4cd7-84d6-612ee3f25051}",
-            "{b54f6be6-0eb5-4a4e-887a-ba9d276f9a83}",
-            "{5ffa3816-358d-4cdd-9b7d-e1f7f5543e05}"
+            "{b020a277-90e2-4cd7-84d6-612ee3f25051}"
         ],
         "value": ""
     },
@@ -770,6 +768,11 @@
         "parent": "data_object",
         "isValue": true,
         "property": "",
+        "dataType": "Float",
+        "association": [
+            "Vertex",
+            "Cell"
+        ],
         "value": 1,
         "min": 1,
         "max": 1000,

simpeg_drivers/potential_fields/gravity/uijson.py

@@ -25,16 +25,55 @@
 from geoh5py.ui_json.ui_json import BaseUIJson
 
 from simpeg_drivers import assets_path
-from simpeg_drivers.uijson import BaseInversionUIJson, SimPEGDriversUIJson
+from simpeg_drivers.uijson import SimPEGDriversUIJson
 
 
 class GravityForwardUIJson(SimPEGDriversUIJson):
     """Gravity Forward UIJson."""
 
-    default_ui_json: ClassVar[Path] = assets_path / "gravity_forward.ui.json"
+    default_ui_json: ClassVar[Path] = assets_path() / "gravity_forward.ui.json"
 
     inversion_type: str
-    forward_only: str
+    forward_only: bool
+    data_object: ObjectForm
+    z_from_topo: BoolForm
+    receivers_radar_drape: DataForm
+    receivers_offset_z: FloatForm
+    gps_receivers_offset: str
+    gz_channel_bool: BoolForm
+    gx_channel_bool: BoolForm
+    gy_channel_bool: BoolForm
+    guv_channel_bool: BoolForm
+    gxy_channel_bool: BoolForm
+    gxx_channel_bool: BoolForm
+    gyy_channel_bool: BoolForm
+    gzz_channel_bool: BoolForm
+    gxz_channel_bool: BoolForm
+    gyz_channel_bool: BoolForm
+    mesh: ObjectForm
+    starting_model: DataForm
+    topography_object: ObjectForm
+    topography: DataForm
+    active_model: DataForm
+    output_tile_files: bool
+    parallelized: BoolForm
+    n_cpu: IntegerForm
+    tile_spatial: DataForm
+    max_chunk_size: IntegerForm
+    chunk_by_rows: BoolForm
+    out_group: GroupForm
+    ga_group: str
+    generate_sweep: BoolForm
+    distributed_workers: str
+
+
+class GravityInversionUIJson(SimPEGDriversUIJson):
+    """Gravity Inversion UIJson."""
+
+    default_ui_json: ClassVar[Path] = assets_path() / "gravity_inversionforward.ui.json"
+
+    inversion_type: str
+    forward_only: bool
     data_object: ObjectForm
     gz_channel: DataForm
     gz_uncertainty: DataForm
@@ -85,34 +124,22 @@ class GravityForwardUIJson(SimPEGDriversUIJson):
     initial_beta: FloatForm
     coolingFactor: FloatForm
     coolingRate: IntegerForm
-    max_global_iteration: IntegerForm
+    max_global_iterations: IntegerForm
     max_line_search_iterations: IntegerForm
     max_cg_iterations: IntegerForm
     tol_cg: FloatForm
     f_min_change: FloatForm
     sens_wts_threshold: FloatForm
     every_iteration_bool: BoolForm
-    save_sensitivity: BoolForm
+    save_sensitivities: BoolForm
     parallelized: BoolForm
     n_cpu: IntegerForm
-    tile_spatial: IntegerForm
+    tile_spatial: DataForm
     store_sensitivities: ChoiceForm
     max_ram: str
     max_chunk_size: IntegerForm
+    chunk_by_rows: BoolForm
     out_group: GroupForm
-
-
-class GravityInversionUIJson(SimPEGDriversUIJson):
-    default_ui_json: ClassVar[Path] = assets_path / "gravity_inversion.ui.json"
-    title: str = "Gravity Inversion"
-    inversion_type: str = "gravity"
-
-    gx_channel_bool: bool = False
-    gy_channel_bool: bool = False
-    gz_channel_bool: bool = True
-    gxx_channel_bool: bool = False
-    gxy_channel_bool: bool = False
-    gxz_channel_bool: bool = False
-    gyy_channel_bool: bool = False
-    gyz_channel_bool: bool = False
-    gzz_channel_bool: bool = False
+    ga_group: str
+    generate_sweep: BoolForm
+    distributed_workers: str

simpeg_drivers/uijson.py

@@ -11,19 +11,27 @@
 import json
 import logging
 from pathlib import Path
-from uuid import UUID
+from typing import Annotated
 
 from geoh5py.groups import SimPEGGroup, UIJsonGroup
+from geoh5py.shared.validators import empty_string_to_none, none_to_empty_string
 from geoh5py.ui_json.enforcers import UUIDEnforcer
 from geoh5py.ui_json.ui_json import BaseUIJson
-from pydantic import field_validator
+from pydantic import BeforeValidator, PlainSerializer, field_validator
 
 import simpeg_drivers
 
 
 logger = logging.getLogger(__name__)
 
 
+OptionalPath = Annotated[
+    Path | None,
+    BeforeValidator(empty_string_to_none),
+    PlainSerializer(none_to_empty_string),
+]
+
+
 class SimPEGDriversUIJson(BaseUIJson):
     version: str = simpeg_drivers.__version__
     title: str
@@ -32,8 +40,8 @@ class SimPEGDriversUIJson(BaseUIJson):
     conda_environment: str
     run_command: str
     geoh5: Path | None
-    monitoring_directory: Path
-    workspace_geoh5: Path
+    monitoring_directory: OptionalPath
+    workspace_geoh5: OptionalPath
 
     @field_validator("version", mode="before")
     @classmethod
@@ -61,80 +69,3 @@ def write_default(cls):
         data = uijson.model_dump_json(indent=4)
         with open(cls.default_ui_json, "w", encoding="utf-8") as file:
             file.write(data)
-
-
-class CoreUIJson(SimPEGDriversUIJson):
-    """
-    Core class for ui.json data.
-    """
-
-    run_command: str
-    conda_environment: str
-    forward_only: bool
-    data_object: UUID
-    mesh: UUID | None
-    starting_model: float | UUID
-    topography_object: UUID | None
-    topography: float | UUID | None
-    active_model: UUID | None
-    tile_spatial: int
-    parallelized: bool
-    n_cpu: int | None
-    max_chunk_size: int
-    save_sensitivities: bool
-    out_group: SimPEGGroup | UIJsonGroup | None
-    generate_sweep: bool
-
-
-class BaseInversionUIJson(BaseUIJson):
-    """
-    Base class for inversion ui.json data.
-    """
-
-    reference_model: float | UUID | None
-    lower_bound: float | UUID | None
-    upper_bound: float | UUID | None
-
-    alpha_s: float | UUID | None
-    length_scale_x: float | UUID
-    length_scale_y: float | UUID
-    length_scale_z: float | UUID
-
-    s_norm: float | UUID | None
-    x_norm: float | UUID
-    y_norm: float | UUID | None
-    z_norm: float | UUID
-    gradient_type: str
-    max_irls_iterations: int
-    starting_chi_factor: float
-
-    chi_factor: float
-    auto_scale_misfits: bool
-    initial_beta_ratio: float | None
-    initial_beta: float | None
-    coolingFactor: float
-
-    coolingRate: float
-    max_global_iterations: int
-    max_line_search_iterations: int
-    max_cg_iterations: int
-    tol_cg: float
-    f_min_change: float
-
-    sens_wts_threshold: float
-    every_iteration_bool: bool
-    save_sensitivities: str
-
-    beta_tol: float
-    prctile: float
-    coolEps_q: bool
-    coolEpsFact: float
-    beta_search: bool
-
-    chunk_by_rows: bool
-    output_tile_files: bool
-    inversion_style: str
-    max_ram: float | None
-    ga_group: SimPEGGroup | None
-    distributed_workers: int | None
-    no_data_value: float | None

tests/uijson_test.py

@@ -13,9 +13,14 @@
 from pathlib import Path
 from typing import ClassVar
 
+import numpy as np
 from geoh5py import Workspace
 
-from simpeg_drivers.uijson import CoreUIJson
+from simpeg_drivers.params import ActiveCellsOptions
+from simpeg_drivers.potential_fields.gravity.params import GravityInversionOptions
+from simpeg_drivers.potential_fields.gravity.uijson import GravityInversionUIJson
+from simpeg_drivers.uijson import SimPEGDriversUIJson
+from simpeg_drivers.utils.testing import setup_inversion_workspace
 
 
 logger = logging.getLogger(__name__)
@@ -25,7 +30,7 @@ def test_version_warning(tmp_path, caplog):
     workspace = Workspace(tmp_path / "test.geoh5")
 
     with caplog.at_level(logging.WARNING):
-        _ = CoreUIJson(
+        _ = SimPEGDriversUIJson(
             version="0.2.0",
             title="My app",
             geoh5=str(workspace.h5file),
@@ -50,7 +55,7 @@ def test_write_default(tmp_path):
     with open(default_path, "w", encoding="utf-8") as f:
         json.dump(data, f, ensure_ascii=False, indent=4)
 
-    class MyUIJson(CoreUIJson):
+    class MyUIJson(SimPEGDriversUIJson):
         default_ui_json: ClassVar[Path] = default_path
         version: str = "0.2.0"
 
@@ -60,3 +65,35 @@ class MyUIJson(CoreUIJson):
         data = json.load(f)
 
     assert data["version"] == "0.3.0-alpha.1"
+
+
+def test_gravity_uijson(tmp_path):
+    geoh5, _, starting_model, survey, topography = setup_inversion_workspace(
+        tmp_path, background=0.0, anomaly=0.75, inversion_type="gravity"
+    )
+
+    opts = GravityInversionOptions(
+        geoh5=geoh5,
+        data_object=survey,
+        gz_channel=survey.add_data({"gz": {"values": np.ones(survey.n_vertices)}}),
+        gz_uncertainty=survey.add_data(
+            {"gz_unc": {"values": np.ones(survey.n_vertices)}}
+        ),
+        mesh=starting_model.parent,
+        starting_model=starting_model,
+        active_cells=ActiveCellsOptions(
+            topography_object=topography,
+        ),
+    )
+    params_uijson_path = tmp_path / "from_params.ui.json"
+    opts.write_ui_json(params_uijson_path)
+
+    uijson = GravityInversionUIJson.read(params_uijson_path)
+    uijson_path = tmp_path / "from_uijson.ui.json"
+    uijson.write(uijson_path)
+    with open(params_uijson_path, encoding="utf-8") as f:
+        params_data = json.load(f)
+    with open(uijson_path, encoding="utf-8") as f:
+        uijson_data = json.load(f)
+
+    assert uijson_data == params_data