From 80726177c9bba2d6f585deec54bfc31384a4b4d3 Mon Sep 17 00:00:00 2001
From: Matt Wittmann <mcwitt@gmail.com>
Date: Fri, 11 Sep 2020 09:42:26 -0700
Subject: [PATCH] Add compound series schema

---
 fah_xchem/schema.py               | 71 +++++++++++++++++++++++++
 fah_xchem/tests/test_fah_xchem.py | 86 ++++++++++++++++++++++++++++++-
 2 files changed, 156 insertions(+), 1 deletion(-)
 create mode 100644 fah_xchem/schema.py

diff --git a/fah_xchem/schema.py b/fah_xchem/schema.py
new file mode 100644
index 00000000..8ebdfc85
--- /dev/null
+++ b/fah_xchem/schema.py
@@ -0,0 +1,71 @@
+import datetime as dt
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel
+
+
+class Model(BaseModel):
+    class Config:
+        allow_mutation = False
+        extra = "forbid"
+
+
+class CompoundSeriesMetadata(Model):
+    name: str
+    description: str
+    creator: str
+    creation_date: dt.date
+    xchem_project: str
+    receptor_variant: Dict[str, str]
+    temperature_kelvin: float
+    ionic_strength_millimolar: float
+    pH: float
+
+
+class Microstate(Model):
+    microstate_id: str
+    smiles: str
+
+
+class Compound(Model):
+    compound_id: str
+    smiles: str
+    experimental_data: Dict[str, float]
+    microstates: List[Microstate]
+
+
+class CompoundMicrostate(Model):
+    compound_id: str
+    microstate_id: str
+
+
+class Transformation(Model):
+    run: int
+    initial_microstate: CompoundMicrostate
+    final_microstate: CompoundMicrostate
+    xchem_fragment_id: str
+
+
+class CompoundSeries(Model):
+    metadata: CompoundSeriesMetadata
+    compounds: List[Compound]
+    transformations: List[Transformation]
+
+
+class ServerConfig(Model):
+    projects_path: str
+    data_path: str
+
+
+class AnalysisConfig(Model):
+    max_binding_delta_f: Optional[float] = None
+    min_num_work_values: Optional[int] = 40
+    work_precision_decimals: Optional[int] = 3
+
+
+class Config(Model):
+    server: ServerConfig
+    complex_project: int
+    solvent_project: int
+    analysis: AnalysisConfig
+    output_dir: str
diff --git a/fah_xchem/tests/test_fah_xchem.py b/fah_xchem/tests/test_fah_xchem.py
index 7194cc5e..20b81522 100644
--- a/fah_xchem/tests/test_fah_xchem.py
+++ b/fah_xchem/tests/test_fah_xchem.py
@@ -3,11 +3,95 @@
 """
 
 # Import package, test suite, and other packages as needed
-import fah_xchem
 import pytest
 import sys
+from fah_xchem.schema import *
+
+
+@pytest.fixture
+def compound_series():
+    return CompoundSeries(
+        metadata=CompoundSeriesMetadata(
+            name="2020-08-20-benzotriazoles",
+            description="Sprint 3: Prioritization of benzotriazole derivatives",
+            creator="John D. Chodera",
+            creation_date=dt.datetime(2020, 9, 8, 10, 14, 48, 607238),
+            xchem_project="Mpro",
+            receptor_variant=dict(
+                biological_assembly="monomer", protein_variant="thiolate"
+            ),
+            temperature_kelvin=300,
+            ionic_strength_millimolar=70,
+            pH=7.4,
+        ),
+        compounds=[
+            Compound(
+                compound_id="MAT-POS-f42f3716-1",
+                smiles="Cc1ccncc1NC(=O)Cc1cc(Cl)cc(-c2ccc(C3CC3(F)F)cc2)c1",
+                experimental_data={"pIC50": 4.324},
+                microstates=[
+                    Microstate(
+                        microstate_id="MAT-POS-f42f3716-1-1",
+                        smiles="Cc1ccncc1NC(=O)Cc1cc(Cl)cc(-c2ccc(C3CC3(F)F)cc2)c1",
+                    ),
+                    Microstate(
+                        microstate_id="MAT-POS-f42f3716-1-2",
+                        smiles="Cc1ccn[H+]cc1NC(=O)Cc1cc(Cl)cc(-c2ccc(C3CC3(F)F)cc2)c1",
+                    ),
+                ],
+            ),
+            Compound(
+                compound_id="MAT-POS-f42f3716-2",
+                smiles="Cc1ccncc1NC(=O)Cc1cc(Cl)cc(-c2ccc(S(C)(=O)=O)cc2Cl)c1",
+                experimental_data={"pIC50": 4.324},
+                microstates=[
+                    Microstate(
+                        microstate_id="MAT-POS-f42f3716-2-1",
+                        smiles="Cc1ccncc1NC(=O)Cc1cc(Cl)cc(-c2ccc(C3CC3(F)F)cc2)c1",
+                    ),
+                    Microstate(
+                        microstate_id="MAT-POS-f42f3716-2-2",
+                        smiles="Cc1ccncc1NC(=O)Cc1cc(Cl)cc(-c2ccc(C3CC3(F)F)cc2)c1",
+                    ),
+                ],
+            ),
+        ],
+        transformations=[
+            Transformation(
+                run=0,
+                initial_microstate=CompoundMicrostate(
+                    compound_id="MAT-POS-f42f3716-1",
+                    microstate_id="MAT-POS-f42f3716-1-1",
+                ),
+                final_microstate=CompoundMicrostate(
+                    compound_id="MAT-POS-f42f3716-1",
+                    microstate_id="MAT-POS-f42f3716-1-2",
+                ),
+                xchem_fragment_id="x10789",
+            ),
+            Transformation(
+                run=1,
+                initial_microstate=CompoundMicrostate(
+                    compound_id="MAT-POS-f42f3716-2",
+                    microstate_id="MAT-POS-f42f3716-2-1",
+                ),
+                final_microstate=CompoundMicrostate(
+                    compound_id="MAT-POS-f42f3716-2",
+                    microstate_id="MAT-POS-f42f3716-2-2",
+                ),
+                xchem_fragment_id="x10789",
+            ),
+        ],
+    )
 
 
 def test_fah_xchem_imported():
     """Sample test, will always pass so long as import statement worked"""
     assert "fah_xchem" in sys.modules
+
+
+def test_compound_series_json_serialization(compound_series):
+    """Test json serialize/deserialize roundtrip for CompoundSeries"""
+    json = compound_series.json()
+    deserialized = CompoundSeries.parse_raw(json)
+    assert deserialized == compound_series