fix: accept equivalent hdf5 psi layouts

Author: edengilbertus

src/io/checkpoint.rs

@@ -166,15 +166,30 @@ pub fn load_checkpoint(path: &Path) -> Result<(Grid3, State3), Error> {
         .dataset(PSI_DATASET)
         .map_err(|_| Error::MissingCheckpointField("state/psi"))?;
     let psi_shape = psi_dataset.shape();
-    if psi_shape.as_slice() != [nx, ny, nz] {
-        return Err(Error::MalformedCheckpointField("state/psi"));
-    }
-    let psi_raw: Vec<ComplexRepr> = psi_dataset
-        .read_raw()
-        .map_err(|_| Error::MalformedCheckpointField("state/psi"))?;
-    if psi_raw.len() != grid.len() {
-        return Err(Error::MalformedCheckpointField("state/psi"));
-    }
+    let psi: Vec<Complex64> = match psi_shape.as_slice() {
+        [sx, sy, sz] if [*sx, *sy, *sz] == [nx, ny, nz] => {
+            let psi_raw: Vec<ComplexRepr> = psi_dataset
+                .read_raw()
+                .map_err(|_| Error::MalformedCheckpointField("state/psi"))?;
+            if psi_raw.len() != grid.len() {
+                return Err(Error::MalformedCheckpointField("state/psi"));
+            }
+            psi_raw.into_iter().map(from_repr).collect()
+        }
+        [sx, sy, sz, 2] if [*sx, *sy, *sz] == [nx, ny, nz] => {
+            let psi_raw: Vec<f64> = psi_dataset
+                .read_raw()
+                .map_err(|_| Error::MalformedCheckpointField("state/psi"))?;
+            if psi_raw.len() != grid.len() * 2 {
+                return Err(Error::MalformedCheckpointField("state/psi"));
+            }
+            psi_raw
+                .chunks_exact(2)
+                .map(|pair| Complex64::new(pair[0], pair[1]))
+                .collect()
+        }
+        _ => return Err(Error::MalformedCheckpointField("state/psi")),
+    };
 
     if let Some(config_group) = maybe_group(&root, CONFIG_GROUP)? {
         validate_optional_group(&config_group, &["dt", "hbar_over_m", "poisson_scale"])?;
@@ -193,10 +208,5 @@ pub fn load_checkpoint(path: &Path) -> Result<(Grid3, State3), Error> {
         )?;
     }
 
-    Ok((
-        grid,
-        State3 {
-            psi: psi_raw.into_iter().map(from_repr).collect(),
-        },
-    ))
+    Ok((grid, State3 { psi }))
 }