Use the eigenvalues to impose the standard Matérn prior for the GP

.gitignore

@@ -25,3 +25,4 @@ mesh_test*
 Untitled*.ipynb
 *.parquet
 local*.ipynb
+**/*.zarr/

celeri/operators.py

@@ -345,6 +345,7 @@ def from_disk(cls, input_dir: str | Path) -> "TdeOperators":
 @dataclass
 class EigenOperators:
     eigenvectors_to_tde_slip: ByMesh[np.ndarray]
+    eigenvalues: ByMesh[np.ndarray]
     eigen_to_velocities: ByMesh[np.ndarray]
     eigen_to_tde_bcs: ByMesh[np.ndarray]
     linear_gaussian_smoothing: ByMesh[np.ndarray]
@@ -585,6 +586,7 @@ class _OperatorBuilder:
     mogi_to_velocities: np.ndarray | None = None
     slip_rate_to_okada_to_velocities: np.ndarray | None = None
     eigenvectors_to_tde_slip: dict[int, np.ndarray] = field(default_factory=dict)
+    eigenvalues: dict[int, np.ndarray] = field(default_factory=dict)
     rotation_to_tri_slip_rate: dict[int, np.ndarray] = field(default_factory=dict)
     linear_gaussian_smoothing: dict[int, np.ndarray] = field(default_factory=dict)
     tde_to_velocities: dict[int, np.ndarray] = field(default_factory=dict)
@@ -651,9 +653,11 @@ def finalize_eigen(self) -> Operators:
         assert self.eigen_to_velocities is not None
         assert self.eigen_to_tde_bcs is not None
         assert self.eigenvectors_to_tde_slip is not None
+        assert self.eigenvalues is not None
 
         eigen = EigenOperators(
             eigenvectors_to_tde_slip=self.eigenvectors_to_tde_slip,
+            eigenvalues=self.eigenvalues,
             linear_gaussian_smoothing=self.linear_gaussian_smoothing,
             eigen_to_velocities=self.eigen_to_velocities,
             eigen_to_tde_bcs=self.eigen_to_tde_bcs,
@@ -2134,15 +2138,18 @@ def _store_eigenvectors_to_tde_slip(model: Model, operators: _OperatorBuilder):
 
     for i in range(len(meshes)):
         logger.info(f"Start: Eigenvectors to TDE slip for mesh: {meshes[i].file_name}")
-        # Get eigenvectors for curren mesh
-        _, eigenvectors = get_eigenvalues_and_eigenvectors(
-            meshes[i].n_modes,
+        # Get eigenvalues and eigenvectors for current mesh
+        eigenvalues, eigenvectors = get_eigenvalues_and_eigenvectors(
+            meshes[i].n_modes_total,  # Use total modes (strike-slip + dip-slip)
             meshes[i].x_centroid,
             meshes[i].y_centroid,
             meshes[i].z_centroid,
             distance_exponent=1.0,  # Make this something set in mesh_parameters.json
         )
 
+        # Store eigenvalues for this mesh
+        operators.eigenvalues[i] = eigenvalues
+
         # Create eigenvectors to TDE slip matrix
         operators.eigenvectors_to_tde_slip[i] = np.zeros(
             (

celeri/solve.py

@@ -585,7 +585,9 @@ def to_disk(self, output_dir: str | Path) -> None:
         self.operators.to_disk(output_dir / "operators")
 
         if self.mcmc_trace is not None:
-            self.mcmc_trace.to_datatree().to_zarr(output_dir / "mcmc_trace.zarr")
+            self.mcmc_trace.to_datatree().to_zarr(
+                output_dir / "mcmc_trace.zarr", mode="w"
+            )
         # We skip saving the trace, it shouldn't be needed later
         dataclass_to_disk(self, output_dir, skip={"operators", "trace", "mcmc_trace"})
 

celeri/solve_mcmc.py

@@ -48,7 +48,7 @@ def _get_eigen_modes(
     operators: Operators,
     out_idx: slice,
 ):
-    """Get the eigenmodes and station velocity operator for a mesh and slip type."""
+    """Get the eigenmodes, eigenvalues and station velocity operator for a mesh and slip type."""
     assert operators.eigen is not None
 
     if kind == "strike_slip":
@@ -64,7 +64,16 @@ def _get_eigen_modes(
     to_velocity = operators.eigen.eigen_to_velocities[mesh][
         :, start_idx : start_idx + n_eigs
     ]
-    return _clean_operator(eigenvectors), _clean_operator(to_velocity)
+    eigenvalues = operators.eigen.eigenvalues[mesh][start_idx : start_idx + n_eigs]
+
+    # Verify that the number of eigenvalues matches the number of eigenvectors
+    if len(eigenvalues) != eigenvectors.shape[1]:
+        raise RuntimeError(
+            f"Eigenvalues and eigenvectors have different numbers of modes. "
+            f"Eigenvalues: {len(eigenvalues)}, eigenvectors: {eigenvectors.shape}"
+        )
+
+    return _clean_operator(eigenvectors), _clean_operator(to_velocity), eigenvalues
 
 
 def _coupling_component(
@@ -97,15 +106,23 @@ def _coupling_component(
     kinematic = operator @ rotation
     pm.Deterministic(f"kinematic_{mesh}_{kind}", kinematic)
 
-    eigenvectors, _ = _get_eigen_modes(
+    eigenvectors, _, eigenvalues = _get_eigen_modes(
         model,
         mesh,
         kind,
         operators,
         out_idx=idx,
     )
     n_eigs = eigenvectors.shape[1]
-    coefs = pm.Normal(f"coupling_coefs_{mesh}_{kind}", mu=0, sigma=10, shape=n_eigs)
+
+    # Use eigenvalue-based variance for Matérn GP prior
+    # Normalize so top eigenmode has std=10
+    eigenvalue_std = np.sqrt(eigenvalues)
+    top_std = eigenvalue_std[0] if len(eigenvalue_std) > 0 else np.nan
+    normalized_std = 10.0 * eigenvalue_std / top_std
+    coefs = pm.Normal(
+        f"coupling_coefs_{mesh}_{kind}", mu=0, sigma=normalized_std, shape=n_eigs
+    )
 
     coupling_field = eigenvectors @ coefs
     coupling_field = _constrain_field(coupling_field, lower, upper)
@@ -147,7 +164,7 @@ def _elastic_component(
     scale = scale / len(operators.eigen.eigen_to_velocities)
     scale = 1 / np.sqrt(scale)
 
-    eigenvectors, to_velocity = _get_eigen_modes(
+    eigenvectors, to_velocity, eigenvalues = _get_eigen_modes(
         model,
         mesh,
         kind,
@@ -156,7 +173,14 @@ def _elastic_component(
     )
     n_eigs = eigenvectors.shape[1]
 
-    raw = pm.Normal(f"elastic_eigen_raw_{mesh}_{kind}", shape=n_eigs)
+    # Use eigenvalue-based variance for Matérn GP prior
+    # Normalize so top eigenmode has std=1
+    eigenvalue_std = np.sqrt(eigenvalues)
+    top_std = eigenvalue_std[0] if len(eigenvalue_std) > 0 else np.nan
+    normalized_std = 1.0 * eigenvalue_std / top_std
+    raw = pm.Normal(
+        f"elastic_eigen_raw_{mesh}_{kind}", sigma=normalized_std, shape=n_eigs
+    )
     param = pm.Deterministic(f"elastic_eigen_{mesh}_{kind}", scale * raw)
     elastic = _constrain_field(eigenvectors @ param, lower, upper)
     pm.Deterministic(f"elastic_{mesh}_{kind}", elastic)