updated

Author: Qwen Code

drift/signaling.py

@@ -111,6 +111,10 @@ def step(self, state, inhibition):
 
     def _custom_step(self, state, inhibition):
         """Step using a custom drift function provided in the topology."""
+        if self.topology.drift_fn is None:
+            # Fallback if drift_fn is None
+            return self._generic_step(state, inhibition)
+
         drift = self.topology.drift_fn(state, self.topology.parameters, inhibition)
         diffusion = np.random.normal(0, self.noise_scale, size=len(state)) * np.sqrt(self.dt)
         new_state = state + drift * self.dt + diffusion

drift/topology.py

@@ -16,11 +16,13 @@ def __init__(
         parameters: Dict[str, float],
         drift_fn: Optional[Callable] = None,
         name: str = "custom_topology",
+        inhibited_species: Optional[str] = None,
     ):
         self.species = species
         self.parameters = parameters
         self.drift_fn = drift_fn
         self.name = name
+        self.inhibited_species = inhibited_species
 
     @classmethod
     def from_json(cls, json_path: str):
@@ -75,13 +77,12 @@ def from_sbml(cls, sbml_path: str, inhibited_species: Optional[str] = None):
         name = model.getName() or model.getId() or "sbml_topology"
         logger.info(f"Imported SBML model: {name}")
 
-        topology = cls(
+        return cls(
             species=species,
             parameters=parameters,
-            name=name
+            name=name,
+            inhibited_species=inhibited_species
         )
-        topology.inhibited_species = inhibited_species
-        return topology
 
     def get_initial_state(self) -> np.ndarray:
         """Returns a default initial state (all 0.5 for normalized)."""