updates to api main and dispatcher

AlexPatrie · AlexPatrie · commit 3360bc818f1c · 2025-02-19T15:57:34.000-05:00
diff --git a/gateway/handlers/states.py b/gateway/handlers/states.py
@@ -91,15 +91,15 @@ def generate_mem3dg_state(
     }
 
     for key, value in parameters_config.items():
-        doc['config']['parameters'][key] = value
+        doc['membrane']['config']['parameters'][key] = value
 
     if geometry_parameters is not None and geometry_type is not None:
-        doc['config']['geometry'] = {
+        doc['membrane']['config']['geometry'] = {
             'type': geometry_type,
             'parameters': geometry_parameters
         }
     else:
-        doc['config']['mesh_file'] = mesh_file
+        doc['membrane']['config']['mesh_file'] = mesh_file
 
     return doc
 
diff --git a/gateway/handlers/submit.py b/gateway/handlers/submit.py
@@ -6,7 +6,7 @@
 from fastapi import UploadFile, HTTPException
 from process_bigraph import Composite
 
-from shared.data_model import UtcRun, AmiciRun, CobraRun, CopasiRun, TelluriumRun, ValidatedComposition
+from shared.data_model import UtcRun, AmiciRun, CobraRun, CopasiRun, TelluriumRun, ValidatedComposition, Mem3dgRun
 from shared.database import DatabaseConnector
 from shared.environment import DEFAULT_JOB_COLLECTION_NAME, DEFAULT_BUCKET_NAME
 from shared.io import write_uploaded_file
@@ -26,6 +26,7 @@ async def submit_pymem3dg_run(
         volume: float,
         parameters_config: dict[str, float | int],
         damping: float,
+        duration: int,
         tolerance: Optional[float] = 1e-11,
         geometry_type: Optional[str] = None,
         geometry_parameters: Optional[Dict[str, Union[float, int]]] = None,
@@ -51,7 +52,9 @@ async def submit_pymem3dg_run(
         job_id=job_id,
         last_updated=db_connector.timestamp(),
         status="PENDING",
-        **input_state
+        simulators=["pymem3dg"],
+        duration=duration,
+        spec=input_state
     )
 
     # save job to db
diff --git a/gateway/main.py b/gateway/main.py
@@ -471,6 +471,7 @@ def check_health() -> HealthCheckResponse:
     summary="Run a Mem3dg Process",
 )
 async def run_mem3dg_process(
+        duration: int = Query(...),
         characteristic_time_step: float = Query(...),
         tension_modulus: float = Query(...),
         preferred_area: float = Query(...),
@@ -479,6 +480,7 @@ async def run_mem3dg_process(
         osmotic_strength: float = Query(...),
         volume: float = Query(...),
         damping: float = Query(...),
+        bending_kbc: float = Query(...),
         tolerance: Optional[float] = Query(default=1e-11),
         # geometry_type: Optional[str] = None,
         # geometry_parameters: Optional[Dict[str, Union[float, int]]] = None,
@@ -491,6 +493,12 @@ async def run_mem3dg_process(
         bucket_name=DEFAULT_BUCKET_NAME,
         extension='.ply'
     )
+
+    parameters = {
+        "bending": {
+            "Kbc": bending_kbc,
+        }
+    }
     mem3dg_run = await submit_pymem3dg_run(
         job_id=job_id,
         db_connector=db_conn_gateway,
@@ -504,6 +512,8 @@ async def run_mem3dg_process(
         damping=damping,
         tolerance=tolerance,
         mesh_file=uploaded_file_location,
+        parameters_config=parameters,
+        duration=duration
     )
 
     return mem3dg_run
diff --git a/pyproject.toml b/pyproject.toml
@@ -20,12 +20,17 @@ dependencies = [
     "nbformat"
 ]
 
+# also make sure to install readdy and pymem3dg:
+# conda install -n bsp conda-forge::pymem3dg -y \
+#     && conda install -n bsp conda-forge::readdy -y
+
 [tool.poetry]
 packages = [
     {include = "gateway"},
     {include = "shared"},
     {include = "worker"},
-    {include = "common"}
+    {include = "common"},
+    {include = "tests"}
 ]
 
 [build-system]
@@ -34,6 +39,7 @@ build-backend = "poetry.core.masonry.api"
 
 
 [tool.pytest.ini_options]
+addopts = "-s"
 testpaths = ["tests"]
-addopts = "--ignore=setup.py"
-python_files = "*.py"
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "function"
diff --git a/shared/data_model.py b/shared/data_model.py
@@ -58,8 +58,20 @@ class SmoldynRun(Run):
     dt: float
 
 
+
+@dataclass
+class Mem3dgRun(CompositionRun):
+    """
+    :param simulators:
+    :param duration:
+    :param spec:
+    :param results:
+    """
+    pass
+
+
 @dataclass
-class Mem3dgRun(Run):
+class _Mem3dgRun(Run):
     characteristic_time_step: float
     tension_modulus: float
     preferred_area: float
diff --git a/shared/log_config.py b/shared/log_config.py
@@ -30,4 +30,5 @@ def setup_logging(fname: str, return_all: bool = False):
     # Add the console handler to the root logger and uvicorn logger
     root_logger.addHandler(console_handler)
 
-    return root_logger, console_handler if return_all else root_logger
+    # return root_logger, console_handler if return_all else root_logger
+    return root_logger
diff --git a/tests/test_dispatcher.py b/tests/test_dispatcher.py
@@ -2,18 +2,61 @@
 import os
 import asyncio
 
+import pytest
+from process_bigraph import Composite
+
 from worker.dispatch import JobDispatcher
 
-TEST_DIR = os.path.dirname(os.path.abspath(__file__))
+# TEST_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+@pytest.fixture
+def test_dir() -> str:
+    return os.path.dirname(os.path.abspath(__file__))
+
+
+@pytest.fixture
+@pytest.mark.usefixtures("test_dir")
+def payload(test_dir: str) -> dict:
+    with open(os.path.join(test_dir, "test_fixtures", "test_membrane_request2.json")) as f:
+        test_request = json.load(f)
+    return test_request
+
 
+@pytest.mark.usefixtures('payload')
+class TestDispatcher:
+    def setup_method(self):
+        self.dispatcher = JobDispatcher()
 
-async def test_dispatch_composition():
-    with open(os.path.join(TEST_DIR, "test_fixtures", "test_request.json")) as f:
+    def test_generate_composition(self, payload: dict):
+        composition = self.dispatcher.generate_composite(payload)
+        assert composition is not None
+        print(f'Got the composition state:\n{composition.state}')
+
+
+def get_test_dir() -> str:
+    return os.path.dirname(os.path.abspath(__file__))
+
+
+def get_payload() -> dict:
+    test_dir = get_test_dir()
+    with open(os.path.join(test_dir, "test_fixtures", "test_membrane_request2.json")) as f:
         test_request = json.load(f)
+    return test_request
 
+
+async def test_generate_composition():
     dispatcher = JobDispatcher()
-    await dispatcher.dispatch_composition(test_request)
+    payload = get_payload()
+    await dispatcher.dispatch_composition(payload)
+    # assert composition is not None
+    # print(f'Got the composition state:\n{composition.state}')
+
+
+asyncio.run(test_generate_composition())
 
+# @pytest.mark.asyncio
+# async def test_dispatch_composition():
+#     dispatcher = JobDispatcher()
+#     await dispatcher.dispatch_composition(test_request)
 
-if __name__ == "__main__":
-    asyncio.run(test_dispatch_composition())
diff --git a/tests/test_fixtures/test_membrane_request.json b/tests/test_fixtures/test_membrane_request.json
@@ -0,0 +1,143 @@
+{
+  "job_id": "composition-9c436329-cede-4680-9628-8483ee6a5528",
+  "last_updated": "2025-02-11 17:36:41.931444",
+  "status": "PENDING",
+  "simulators": [],
+  "duration": 10,
+  "spec": {
+    "membrane": {
+      "_type": "process",
+      "address": "local:simple-membrane-process",
+      "config": {
+        "characteristic_time_step": 1,
+        "geometry": {
+          "type": "icosphere",
+          "parameters": {
+            "radius": 1,
+            "subdivision": 2
+          }
+        },
+        "tension_model": {
+          "modulus": 0.1,
+          "preferred_area": 12.4866
+        },
+        "osmotic_model": {
+          "preferred_volume": 2.9306666666666668,
+          "reservoir_volume": 1,
+          "strength": 0.02,
+          "volume": 2.9
+        },
+        "parameters": {
+          "bending": {
+            "Kbc": 0.0000822
+          },
+          "damping": 0.05
+        },
+        "tolerance": 1e-11,
+        "console_output": false
+      },
+      "inputs": {
+        "geometry": [
+          "geometry_store"
+        ],
+        "velocities": [
+          "velocities_store"
+        ],
+        "protein_density": [
+          "protein_density_store"
+        ],
+        "volume": [
+          "volume_store"
+        ],
+        "preferred_volume": [
+          "preferred_volume_store"
+        ],
+        "reservoir_volume": [
+          "reservoir_volume_store"
+        ],
+        "surface_area": [
+          "surface_area_store"
+        ],
+        "osmotic_strength": [
+          "osmotic_strength_store"
+        ]
+      },
+      "outputs": {
+        "geometry": [
+          "geometry_store"
+        ],
+        "velocities": [
+          "velocities_store"
+        ],
+        "protein_density": [
+          "protein_density_store"
+        ],
+        "volume": [
+          "volume_store"
+        ],
+        "preferred_volume": [
+          "preferred_volume_store"
+        ],
+        "reservoir_volume": [
+          "reservoir_volume_store"
+        ],
+        "surface_area": [
+          "surface_area_store"
+        ],
+        "net_forces": [
+          "net_forces_store"
+        ],
+        "notable_vertices": [
+          "notable_vertices_store"
+        ]
+      }
+    },
+    "emitter": {
+      "_type": "step",
+      "address": "local:ram-emitter",
+      "config": {
+        "emit": {
+          "geometry": "GeometryType",
+          "velocities": "VelocitiesType",
+          "protein_density": "ProteinDensityType",
+          "volume": "float",
+          "preferred_volume": "float",
+          "reservoir_volume": "float",
+          "surface_area": "float",
+          "net_forces": "MechanicalForcesType",
+          "notable_vertices": "list[boolean]"
+        }
+      },
+      "inputs": {
+        "geometry": [
+          "geometry_store"
+        ],
+        "velocities": [
+          "velocities_store"
+        ],
+        "protein_density": [
+          "protein_density_store"
+        ],
+        "volume": [
+          "volume_store"
+        ],
+        "preferred_volume": [
+          "preferred_volume_store"
+        ],
+        "reservoir_volume": [
+          "reservoir_volume_store"
+        ],
+        "surface_area": [
+          "surface_area_store"
+        ],
+        "net_forces": [
+          "net_forces_store"
+        ],
+        "notable_vertices": [
+          "notable_vertices_store"
+        ]
+      }
+    }
+  },
+  "results": {}
+}
diff --git a/tests/test_fixtures/test_membrane_request2.json b/tests/test_fixtures/test_membrane_request2.json
diff --git a/tests/test_worker.py b/tests/test_worker.py
diff --git a/worker/dispatch.py b/worker/dispatch.py
