Preserve NWB subject age reference in metadata

Author: AtomicGlance

dandi/consts.py

@@ -32,6 +32,7 @@
 
 metadata_nwb_subject_fields = (
     "age",
+    "age__reference",
     "date_of_birth",
     "genotype",
     "sex",

dandi/metadata/util.py

@@ -247,6 +247,8 @@ def extract_age(metadata: dict) -> models.PropertyValue | None:
     except (KeyError, TypeError, ValueError):
         if metadata.get("age") is not None:
             duration, ref = parse_age(metadata["age"])
+            if metadata.get("age__reference") == "gestational":
+                ref = "Gestational"
         else:
             return None
     else:

dandi/tests/test_metadata.py

@@ -34,6 +34,7 @@
 import numpy as np
 from pydantic import ByteSize
 from pynwb import NWBHDF5IO, NWBFile, TimeSeries
+from pynwb.file import Subject
 import pytest
 import requests
 from semantic_version import Version
@@ -475,6 +476,59 @@ def test_time_extract() -> None:
     )
 
 
+@pytest.mark.ai_generated
+@pytest.mark.parametrize("reference", ["birth", "gestational"])
+def test_age_reference_roundtrip(tmp_path: Path, reference: str) -> None:
+    path = tmp_path / "age.nwb"
+    nwb = NWBFile("age reference", "age", datetime(2020, 1, 1, tzinfo=tzutc()))
+    nwb.subject = Subject(age="P3W", age__reference=reference)
+    with NWBHDF5IO(str(path), "w") as io:
+        io.write(nwb)
+    metadata = get_metadata(path)
+    assert metadata["age__reference"] == reference
+    age = extract_age(metadata)
+    assert age is not None
+    assert age.value == "P3W"
+    expected = (
+        AgeReferenceType.GestationalReference
+        if reference == "gestational"
+        else AgeReferenceType.BirthReference
+    )
+    assert age.valueReference == PropertyValue(value=expected)
+
+
+@pytest.mark.ai_generated
+@pytest.mark.parametrize("reference", ["birth", None])
+@pytest.mark.parametrize("legacy", [False, True])
+def test_age_reference_legacy(reference: str | None, legacy: bool) -> None:
+    age = extract_age(
+        {"age": "Gestational week 3" if legacy else "P3W", "age__reference": reference}
+    )
+    assert age is not None
+    assert age.value == "P3W"
+    assert age.valueReference == PropertyValue(
+        value=AgeReferenceType.GestationalReference
+        if legacy
+        else AgeReferenceType.BirthReference
+    )
+
+
+@pytest.mark.ai_generated
+@pytest.mark.parametrize("reference", ["birth", "gestational", None])
+def test_age_reference_preserves_birth_dates(reference: str | None) -> None:
+    age = extract_age(
+        {
+            "age": "P3W",
+            "age__reference": reference,
+            "date_of_birth": "2020-01-01T00:00:00Z",
+            "session_start_time": "2020-01-03T00:00:00Z",
+        }
+    )
+    assert age is not None
+    assert age.value == "P2D"
+    assert age.valueReference == PropertyValue(value=AgeReferenceType.BirthReference)
+
+
 def test_time_extract_gest() -> None:
     """extract age with Gestational ref"""
     meta_birth = {