Merge pull request #409 from macrocosm-os/staging

Staging

Author: Sarkosos

folding/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.3.4"
+__version__ = "2.3.5"
 version_split = __version__.split(".")
 __spec_version__ = (
     (10000 * int(version_split[0]))

folding/miners/folding_miner.py

@@ -20,7 +20,11 @@
 from folding.base.miner import BaseMinerNeuron
 from folding.base.simulation import OpenMMSimulation
 from folding.protocol import JobSubmissionSynapse
-from folding.utils.reporters import ExitFileReporter, LastTwoCheckpointsReporter
+from folding.utils.reporters import (
+    ExitFileReporter,
+    LastTwoCheckpointsReporter,
+    ProteinStructureReporter,
+)
 from folding.utils.ops import (
     check_if_directory_exists,
     get_tracebacks,
@@ -941,21 +945,24 @@ def configure_commands(
                     reportInterval=self.CHECKPOINT_INTERVAL,
                 )
             )
-            simulation.reporters.append(
-                app.StateDataReporter(
-                    file=f"{self.output_dir}/{state}.log",
-                    reportInterval=self.STATE_DATA_REPORTER_INTERVAL,
-                    step=True,
-                    potentialEnergy=True,
-                )
-            )
+
             simulation.reporters.append(
                 ExitFileReporter(
                     filename=f"{self.output_dir}/{state}",
                     reportInterval=self.EXIT_REPORTER_INTERVAL,
                     file_prefix=state,
                 )
             )
+            simulation.reporters.append(
+                ProteinStructureReporter(
+                    file=f"{self.output_dir}/{state}.log",
+                    reportInterval=self.STATE_DATA_REPORTER_INTERVAL,
+                    step=True,
+                    potentialEnergy=True,
+                    reference_pdb=os.path.join(self.output_dir, f"{self.pdb_id}.pdb"),
+                    speed=True,
+                )
+            )
             state_commands[state] = simulation
 
         return state_commands

folding/organic/organic.py

@@ -1,7 +1,9 @@
+from typing import Optional
 import uuid
 import time
 import json
-from fastapi import APIRouter, Request, Depends, HTTPException
+import tempfile
+from fastapi import APIRouter, Request, Depends, HTTPException, UploadFile, File, Form
 from folding_api.schemas import EpistulaHeaders
 from folding_api.schemas import FoldingParams
 from folding.utils.logging import logger
@@ -10,22 +12,14 @@
 router = APIRouter()
 
 
-@router.post("/organic")
-async def organic(
+def verify_organic_request(
     request: Request,
     job: FoldingParams,
-    epistula_headers: EpistulaHeaders = Depends(EpistulaHeaders),
-):
+    epistula_headers: EpistulaHeaders,
+) -> None:
     """
-    This endpoint is used to receive organic requests. Returns success message with the job id.
-    Args:
-        request: Request
-        job: FoldingParams
-        epistula_headers: EpistulaHeaders
-    Returns:
-        dict[str, str]: dict with the job id.
+    Verify the organic request signature and whitelist.
     """
-
     body_bytes = json.dumps(job.model_dump(), default=str, sort_keys=True).encode(
         "utf-8"
     )
@@ -37,7 +31,6 @@ async def organic(
         raise HTTPException(status_code=403, detail=str(e))
 
     sender_hotkey = epistula_headers.signed_by
-
     if sender_hotkey not in request.app.state.config.organic_whitelist:
         logger.warning(
             f"Received organic request from {sender_hotkey}, but {sender_hotkey} is not in the whitelist."
@@ -46,9 +39,87 @@ async def organic(
             status_code=403, detail="Forbidden, sender not in whitelist."
         )
 
+
+def get_folding_params(query: str = Form(...)) -> FoldingParams:
+    """
+    Dependency function to parse and validate the query form data.
+    """
+    try:
+        query_data = json.loads(query)
+        return FoldingParams(**query_data)
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Invalid query data: {str(e)}")
+
+
+@router.post("/organic")
+async def organic(
+    request: Request,
+    job: FoldingParams,
+    epistula_headers: EpistulaHeaders = Depends(EpistulaHeaders),
+):
+    """
+    This endpoint is used to receive organic requests for proteins from RCSB or PDBE databases.
+    Returns success message with the job id.
+
+    Args:
+        request: Request
+        job: FoldingParams
+        epistula_headers: EpistulaHeaders
+    Returns:
+        dict[str, str]: dict with the job id.
+    """
+    verify_organic_request(request, job, epistula_headers)
+
     folding_params = job.model_dump()
     folding_params["job_id"] = str(uuid.uuid4())
+
     logger.info(f"Received organic request: {folding_params}")
     request.app.state.validator._organic_queue.add(folding_params)
 
     return {"job_id": folding_params["job_id"]}
+
+
+@router.post("/organic/upload")
+async def organic_with_upload(
+    request: Request,
+    job: FoldingParams = Depends(get_folding_params),
+    pdb_file: UploadFile = File(...),
+    epistula_headers: EpistulaHeaders = Depends(EpistulaHeaders),
+):
+    """
+    This endpoint is used to receive organic requests with custom PDB files.
+    Returns success message with the job id.
+
+    Args:
+        request: Request
+        job: FoldingParams
+        pdb_file: PDB file upload
+        epistula_headers: EpistulaHeaders
+    Returns:
+        dict[str, str]: dict with the job id.
+    """
+    verify_organic_request(request, job, epistula_headers)
+
+    folding_params = job.model_dump()
+    folding_params["job_id"] = str(uuid.uuid4())
+
+    # Handle PDB file
+    try:
+        # Create a temporary file to store the PDB
+        with tempfile.NamedTemporaryFile(
+            mode="wb", suffix=".pdb", delete=False
+        ) as temp_pdb:
+            content = await pdb_file.read()
+            temp_pdb.write(content)
+            temp_pdb_path = temp_pdb.name
+
+        # Update folding params with the temporary file path
+        folding_params["pdb_file_path"] = temp_pdb_path
+        logger.info(f"Created temporary PDB file at: {temp_pdb_path}")
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Invalid PDB file: {str(e)}")
+
+    logger.info(f"Received organic request with PDB file: {folding_params}")
+    request.app.state.validator._organic_queue.add(folding_params)
+
+    return {"job_id": folding_params["job_id"]}

folding/store.py

@@ -276,7 +276,7 @@ def upload_job(
             s3_links=event["s3_links"],
             priority=event.get("priority", 1),
             update_interval=event.get(
-                "update_interval", random.randint(7200, 14400)
+                "time_to_live", random.randint(7200, 14400)
             ),  # between 2 hours and 4 hours in seconds
             max_time_no_improvement=event.get("max_time_no_improvement", 1),
             is_organic=event.get("is_organic", False),

folding/utils/config.py

@@ -366,7 +366,7 @@ def add_validator_args(cls, parser):
         "--neuron.vpermit_tao_limit",
         type=int,
         help="The maximum number of TAO allowed to query a validator with a vpermit.",
-        default=4096,
+        default=20_000,
     )
 
     parser.add_argument(

folding/utils/reporters.py

@@ -1,5 +1,12 @@
+import io
 import os
+import time
 import openmm.app as app
+import MDAnalysis as mda
+from MDAnalysis.analysis.rms import rmsd
+import numpy as np
+
+
 
 
 class LastTwoCheckpointsReporter(app.CheckpointReporter):
@@ -38,3 +45,135 @@ def report(self, simulation, state):
 
     def finalize(self):
         pass
+
+
+class ProteinStructureReporter(app.StateDataReporter):
+    def __init__(
+        self, file, reportInterval, reference_pdb, **kwargs
+    ):
+        super().__init__(file, reportInterval, **kwargs)
+        self.reference_universe = mda.Universe(reference_pdb)
+        self.positions_history = []  # Store positions for RMSF calculation
+
+    def report(self, simulation, state):
+        """Generate a report.
+
+        Parameters
+        ----------
+        simulation : Simulation
+            The Simulation to generate a report for
+        state : State
+            The current state of the simulation
+        """
+        if not self._hasInitialized:
+            self._initializeConstants(simulation)
+            headers = self._constructHeaders()
+            if not self._append:
+                print(
+                    '#"%s"' % ('"' + self._separator + '"').join(headers),
+                    file=self._out,
+                )
+            try:
+                self._out.flush()
+            except AttributeError:
+                pass
+            self._initialClockTime = time.time()
+            self._initialSimulationTime = state.getTime()
+            self._initialSteps = simulation.currentStep
+            self._hasInitialized = True
+
+        # Check for errors.
+        self._checkForErrors(simulation, state)
+
+        # Store current positions for RMSF calculation
+        universe = self.create_mda_universe(simulation)
+        self.positions_history.append(universe.select_atoms("backbone").positions.copy())
+
+        # Query for the values
+        values = self._constructReportValues(simulation, state)
+
+        # Write the values.
+        print(self._separator.join(str(v) for v in values), file=self._out)
+        try:
+            self._out.flush()
+        except AttributeError:
+            pass
+
+    def _constructReportValues(self, simulation, state):
+        values = super()._constructReportValues(simulation, state)
+        rmsd = self._calculate_rmsd(self.create_mda_universe(simulation))
+        rmsf = self._calculate_rmsf()
+        values.extend([rmsd, rmsf])
+        return values
+
+    def _calculate_rmsd(self, universe):
+        """Calculate RMSD between current and reference positions.
+
+        Args:
+            positions (np.ndarray): Current positions
+
+        Returns:
+            float: RMSD
+        """
+        current_positions = universe.select_atoms("backbone").positions.copy()
+        reference_positions = self.reference_universe.select_atoms("backbone").positions.copy()
+        rmsd_measure = rmsd(current_positions, reference_positions, center=True)
+        return rmsd_measure
+
+    def _calculate_rmsf(self):
+        """Calculate RMSF (Root Mean Square Fluctuation) over time.
+
+        Returns:
+            float: RMSF value in nanometers
+        """
+        if len(self.positions_history) < 2:
+            return 0.0
+            
+        # Convert positions history to numpy array
+        positions_array = np.array(self.positions_history)
+        
+        # Calculate mean position for each atom
+        mean_positions = np.mean(positions_array, axis=0)
+        
+        # Calculate RMSF
+        squared_diff = np.square(positions_array - mean_positions)
+        rmsf = np.sqrt(np.mean(squared_diff))
+        
+        # Keep only the last 1000 frames to prevent memory issues
+        if len(self.positions_history) > 1000:
+            self.positions_history = self.positions_history[-1000:]
+            
+        return rmsf
+
+    def _constructHeaders(self):
+        headers = super()._constructHeaders()
+        headers.extend(["RMSD", "RMSF"])
+        return headers
+
+    def create_mda_universe(self,simulation):
+        """
+        Create an MDAnalysis Universe from an OpenMM simulation object.
+        
+        Args:
+            simulation (openmm.app.Simulation): The OpenMM simulation object
+            
+        Returns:
+            mda.Universe: An MDAnalysis Universe containing the current state of the simulation
+        """
+        # Get the current state
+        state = simulation.context.getState(getPositions=True)
+        positions = state.getPositions(asNumpy=True)
+        
+        # Get the topology
+        topology = simulation.topology
+        
+        # Create a PDB string from the current state
+        pdb_string = io.StringIO()
+        app.PDBFile.writeFile(topology, positions, pdb_string)
+        pdb_string.seek(0)
+        
+        # Create MDAnalysis Universe from the PDB string
+        universe = mda.Universe(pdb_string, format='pdb')
+        
+        return universe
+   

folding/utils/uids.py

@@ -65,7 +65,9 @@ def get_random_uids(self, k: int, exclude: List[int] = None) -> torch.LongTensor
     return uids
 
 
-def get_all_miner_uids(metagraph, vpermit_tao_limit, include_serving_in_check: bool = True) -> List[int]:
+def get_all_miner_uids(
+    metagraph, vpermit_tao_limit, include_serving_in_check: bool = True
+) -> List[int]:
     """Returns all available miner uids from the metagraph.
     Returns:
         uids (List): All available miner uids.
@@ -74,10 +76,10 @@ def get_all_miner_uids(metagraph, vpermit_tao_limit, include_serving_in_check: b
     candidate_uids = []
     for uid in range(metagraph.n.item()):
         uid_is_available = check_uid_availability(
-            metagraph = metagraph,
-            uid = uid,
-            vpermit_tao_limit = vpermit_tao_limit,
-            include_serving_in_check = include_serving_in_check,
+            metagraph=metagraph,
+            uid=uid,
+            vpermit_tao_limit=vpermit_tao_limit,
+            include_serving_in_check=include_serving_in_check,
         )
         if uid_is_available:
             candidate_uids.append(uid)