neorx 0.2.0__py3-none-any.whl

neorx/__init__.py

@@ -0,0 +1,118 @@
+"""
+NeoRx — Public API
+===========================
+
+Thin re-export package so users can write clean imports::
+
+    from neorx import run_pipeline, build_disease_graph
+    from neorx import identify_causal_targets
+    from neorx import predict_admet, save_graph
+
+Instead of the internal layout::
+
+    from neorx.core import run_pipeline   # also works
+"""
+
+from neorx.core import (  # noqa: F401
+    # Enums
+    NodeType,
+    EdgeType,
+    JobStatus,
+    TargetClassification,
+    # Graph models
+    GraphNode,
+    GraphEdge,
+    DiseaseGraph,
+    # Analysis models
+    NeoRxResult,
+    ScoredCandidate,
+    # Pipeline models
+    PipelineJob,
+    PipelineResult,
+    # API models
+    RunRequest,
+    GraphRequest,
+    IdentifyRequest,
+    ScreenTargetRequest,
+    StatusResponse,
+    # Core functions
+    build_disease_graph,
+    disease_graph_to_networkx,
+    identify_causal_targets,
+    score_candidate,
+    rank_candidates,
+    normalise_affinity,
+    normalise_sa,
+    run_pipeline,
+    generate_report,
+    # Cache
+    get_cache,
+    cached_api_call,
+    store_api_response,
+    # Persistence
+    save_graph,
+    load_graph,
+    list_saved_graphs,
+    # ADMET
+    predict_admet,
+    ADMETProfile,
+)
+
+# --- Subpackages -----------------------------------------------------
+from neorx import causalbiorl, core, dockbot, genmol, mirrorfold, molscreen  # noqa: F401,E402
+
+__version__ = "0.2.0"
+
+__all__ = [
+    # Enums
+    "NodeType",
+    "EdgeType",
+    "JobStatus",
+    "TargetClassification",
+    # Graph models
+    "GraphNode",
+    "GraphEdge",
+    "DiseaseGraph",
+    # Analysis models
+    "NeoRxResult",
+    "ScoredCandidate",
+    # Pipeline models
+    "PipelineJob",
+    "PipelineResult",
+    # API models
+    "RunRequest",
+    "GraphRequest",
+    "IdentifyRequest",
+    "ScreenTargetRequest",
+    "StatusResponse",
+    # Core functions
+    "build_disease_graph",
+    "disease_graph_to_networkx",
+    "identify_causal_targets",
+    "score_candidate",
+    "rank_candidates",
+    "normalise_affinity",
+    "normalise_sa",
+    "run_pipeline",
+    "generate_report",
+    # Cache
+    "get_cache",
+    "cached_api_call",
+    "store_api_response",
+    # Persistence
+    "save_graph",
+    "load_graph",
+    "list_saved_graphs",
+    # ADMET
+    "predict_admet",
+    "ADMETProfile",
+]
+
+__all__ = __all__ + [
+    "core",
+    "genmol",
+    "causalbiorl",
+    "molscreen",
+    "dockbot",
+    "mirrorfold",
+]

neorx/causalbiorl/__init__.py

@@ -0,0 +1,23 @@
+"""
+CausalBioRL — Causal Reinforcement Learning Environments for Biological System Control.
+
+Provides Gymnasium-compatible RL environments that simulate biological systems
+(gene expression, metabolic flux, cell populations) and causal RL agents that
+use structural causal models as world models.
+
+Reference:
+    CausalBioRL: Causal Reinforcement Learning Environments for
+    Biological System Control (2026).
+"""
+
+from neorx.causalbiorl.envs.registration import register_envs
+
+from . import agents, causal, envs  # noqa: F401,E402
+from .agents.causal_agent import CausalAgent  # noqa: F401,E402
+from .envs.drug_discovery import DrugDiscoveryEnv  # noqa: F401,E402
+
+__version__ = "0.1.0"
+__all__ = ["envs", "agents", "causal", "DrugDiscoveryEnv", "CausalAgent"]
+
+# Register Gymnasium environments on import
+register_envs()

neorx/causalbiorl/__main__.py

@@ -0,0 +1,193 @@
+"""
+CausalBioRL command-line interface.
+
+Usage:
+    python -m causalbiorl train --env GeneticToggle-v0 --agent causal --episodes 1000
+    python -m causalbiorl benchmark --envs all --agents all --seeds 10
+    python -m causalbiorl play --env GeneticToggle-v0
+    python -m causalbiorl visualise --env GeneticToggle-v0 --agent causal
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Annotated, Optional
+
+import gymnasium as gym
+import numpy as np
+import typer
+
+import neorx.causalbiorl  # register envs  # noqa: F401
+
+app = typer.Typer(
+    name="causalbiorl",
+    help="CausalBioRL — Causal RL environments for biological system control.",
+    add_completion=False,
+)
+
+
+# ────────────────────────────────────────────────────────────────────────── #
+#  train                                                                     #
+# ────────────────────────────────────────────────────────────────────────── #
+
+
+@app.command()
+def train(
+    env: Annotated[str, typer.Option(help="Gymnasium environment ID")] = "GeneticToggle-v0",
+    agent: Annotated[str, typer.Option(help="Agent type: causal, ppo, sac, random")] = "causal",
+    episodes: Annotated[int, typer.Option(help="Number of training episodes")] = 500,
+    difficulty: Annotated[str, typer.Option(help="easy / medium / hard")] = "medium",
+    seed: Annotated[Optional[int], typer.Option(help="Random seed")] = None,
+    output_dir: Annotated[str, typer.Option(help="Output directory")] = "results",
+) -> None:
+    """Train an agent on a CausalBioRL environment."""
+    from neorx.causalbiorl.benchmark import _make_agent
+
+    typer.echo(f"🧬 Training {agent} on {env} ({difficulty}) for {episodes} episodes …")
+    environment = gym.make(env, difficulty=difficulty)
+    ag = _make_agent(agent, environment, seed=seed or 0)
+    metrics = ag.train(n_episodes=episodes, verbose=True)
+    environment.close()
+
+    out = Path(output_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    # Save reward curve
+    import json
+    with open(out / f"{env}_{agent}_rewards.json", "w") as f:
+        json.dump(metrics["episode_rewards"], f)
+
+    final_mean = float(np.mean(metrics["episode_rewards"][-50:]))
+    typer.echo(f"✅ Done — final 50-episode mean reward: {final_mean:.3f}")
+
+
+# ────────────────────────────────────────────────────────────────────────── #
+#  benchmark                                                                 #
+# ────────────────────────────────────────────────────────────────────────── #
+
+
+@app.command()
+def benchmark(
+    envs: Annotated[str, typer.Option(help="Comma-separated env IDs or 'all'")] = "all",
+    agents: Annotated[str, typer.Option(help="Comma-separated agent types or 'all'")] = "all",
+    seeds: Annotated[int, typer.Option(help="Number of random seeds")] = 10,
+    episodes: Annotated[int, typer.Option(help="Episodes per run")] = 500,
+    difficulty: Annotated[str, typer.Option(help="easy / medium / hard")] = "medium",
+    output_dir: Annotated[str, typer.Option(help="Output directory")] = "results",
+) -> None:
+    """Run the full benchmark suite."""
+    from neorx.causalbiorl.benchmark import ALL_AGENTS, ALL_ENVS, run_benchmark
+
+    env_list = ALL_ENVS if envs == "all" else [e.strip() for e in envs.split(",")]
+    agent_list = ALL_AGENTS if agents == "all" else [a.strip() for a in agents.split(",")]
+
+    typer.echo(f"🧬 Benchmark: {env_list} × {agent_list} × {seeds} seeds")
+    run_benchmark(
+        envs=env_list,
+        agents=agent_list,
+        n_seeds=seeds,
+        n_episodes=episodes,
+        difficulty=difficulty,
+        output_dir=output_dir,
+    )
+    typer.echo("✅ Benchmark complete — results saved to " + output_dir)
+
+
+# ────────────────────────────────────────────────────────────────────────── #
+#  play                                                                      #
+# ────────────────────────────────────────────────────────────────────────── #
+
+
+@app.command()
+def play(
+    env: Annotated[str, typer.Option(help="Gymnasium environment ID")] = "GeneticToggle-v0",
+    difficulty: Annotated[str, typer.Option(help="easy / medium / hard")] = "easy",
+    steps: Annotated[int, typer.Option(help="Max steps")] = 100,
+) -> None:
+    """Interactive manual control of an environment.
+
+    Each step prompts you for action values (comma-separated floats in [0, 1]).
+    """
+    environment = gym.make(env, difficulty=difficulty)
+    state, info = environment.reset()
+    action_dim = int(np.prod(environment.action_space.shape))  # type: ignore[union-attr]
+
+    typer.echo(f"🎮 Playing {env} ({difficulty}) — action dim = {action_dim}")
+    typer.echo(f"   State: {state}")
+
+    for t in range(steps):
+        raw = typer.prompt(f"Step {t} — enter {action_dim} action values (comma-sep, 0-1)")
+        try:
+            action = np.array([float(x.strip()) for x in raw.split(",")], dtype=np.float32)
+        except ValueError:
+            typer.echo("  ⚠ Invalid input — using zeros.")
+            action = np.zeros(action_dim, dtype=np.float32)
+
+        state, reward, terminated, truncated, info = environment.step(action)
+        typer.echo(f"   State: {np.round(state, 3)}  |  Reward: {reward:.3f}")
+
+        if terminated or truncated:
+            typer.echo("  Episode finished.")
+            break
+
+    environment.close()
+
+
+# ────────────────────────────────────────────────────────────────────────── #
+#  visualise                                                                 #
+# ────────────────────────────────────────────────────────────────────────── #
+
+
+@app.command()
+def visualise(
+    env: Annotated[str, typer.Option(help="Gymnasium environment ID")] = "GeneticToggle-v0",
+    agent: Annotated[str, typer.Option(help="Agent type")] = "causal",
+    episodes: Annotated[int, typer.Option(help="Training episodes before viz")] = 200,
+    difficulty: Annotated[str, typer.Option(help="Difficulty")] = "medium",
+    output_dir: Annotated[str, typer.Option(help="Output directory")] = "results",
+) -> None:
+    """Train an agent and visualise the learned causal graph + trajectories."""
+    from neorx.causalbiorl.benchmark import _make_agent
+    from neorx.causalbiorl.viz import plot_causal_graph, plot_trajectories
+
+    typer.echo(f"🔬 Training {agent} on {env} for visualisation …")
+    environment = gym.make(env, difficulty=difficulty)
+    ag = _make_agent(agent, environment, seed=42)
+    metrics = ag.train(n_episodes=episodes, verbose=True)
+    environment.close()
+
+    out = Path(output_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    # Reward trajectory
+    plot_trajectories(
+        {agent: metrics["episode_rewards"]},
+        title=f"{env} — {agent}",
+        save_path=out / f"{env}_{agent}_trajectory.png",
+    )
+
+    # Causal graph (if causal agent)
+    if hasattr(ag, "get_learned_graph"):
+        graph = ag.get_learned_graph()
+        if graph is not None:
+            plot_causal_graph(
+                graph,
+                title=f"Learned Causal Graph — {env}",
+                save_path=out / f"{env}_{agent}_causal_graph.png",
+            )
+            typer.echo(f"  Learned graph: {graph.number_of_nodes()} nodes, {graph.number_of_edges()} edges")
+
+    typer.echo(f"✅ Visualisations saved to {out}")
+
+
+# ────────────────────────────────────────────────────────────────────────── #
+#  Entry point                                                               #
+# ────────────────────────────────────────────────────────────────────────── #
+
+
+def main() -> None:
+    app()
+
+
+if __name__ == "__main__":
+    main()

neorx/causalbiorl/agents/__init__.py

@@ -0,0 +1,6 @@
+"""CausalBioRL agents — causal and baseline RL agents."""
+
+from neorx.causalbiorl.agents.causal_agent import CausalAgent
+from neorx.causalbiorl.agents.baseline_agent import PPOAgent, SACAgent, RandomAgent
+
+__all__ = ["CausalAgent", "PPOAgent", "SACAgent", "RandomAgent"]

neorx/causalbiorl/agents/baseline_agent.py

@@ -0,0 +1,161 @@
+"""
+Baseline RL agents for comparison — PPO, SAC, and Random.
+
+Wraps ``stable-baselines3`` implementations with a consistent interface
+so that benchmarking code can treat all agents uniformly.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+from numpy.typing import NDArray
+from tqdm import tqdm
+
+
+# ────────────────────────────────────────────────────────────────────────── #
+#  Random agent                                                              #
+# ────────────────────────────────────────────────────────────────────────── #
+
+
+class RandomAgent:
+    """Agent that samples uniformly from the action space."""
+
+    def __init__(self, env: gym.Env, seed: int | None = None) -> None:
+        self.env = env
+        self.rng = np.random.default_rng(seed)
+        self.episode_rewards: list[float] = []
+        self.episode_lengths: list[int] = []
+
+    def train(
+        self, n_episodes: int = 500, verbose: bool = True, **_kwargs: object
+    ) -> dict[str, Any]:
+        pbar = tqdm(range(n_episodes), desc="RandomAgent", disable=not verbose)
+        total_steps = 0
+        for _ in pbar:
+            state, _ = self.env.reset(seed=int(self.rng.integers(0, 2**31)))
+            ep_reward = 0.0
+            ep_len = 0
+            terminated = truncated = False
+            while not (terminated or truncated):
+                action = self.env.action_space.sample()
+                state, reward, terminated, truncated, _ = self.env.step(action)
+                ep_reward += float(reward)
+                ep_len += 1
+                total_steps += 1
+            self.episode_rewards.append(ep_reward)
+            self.episode_lengths.append(ep_len)
+        return {
+            "episode_rewards": self.episode_rewards,
+            "episode_lengths": self.episode_lengths,
+            "total_steps": total_steps,
+        }
+
+    def act(self, state: NDArray[np.floating]) -> NDArray[np.floating]:
+        return self.env.action_space.sample()
+
+
+# ────────────────────────────────────────────────────────────────────────── #
+#  Stable-Baselines3 wrappers                                                #
+# ────────────────────────────────────────────────────────────────────────── #
+
+
+class _SB3Agent:
+    """Base wrapper for stable-baselines3 algorithms."""
+
+    _algo_cls: type | None = None
+
+    def __init__(
+        self,
+        env: gym.Env,
+        seed: int | None = None,
+        policy: str = "MlpPolicy",
+        **sb3_kwargs: object,
+    ) -> None:
+        self.env = env
+        self.seed = seed
+        self.episode_rewards: list[float] = []
+        self.episode_lengths: list[int] = []
+
+        try:
+            import stable_baselines3  # noqa: F401
+        except ImportError as exc:
+            raise ImportError(
+                "stable-baselines3 is required for PPO/SAC baselines: "
+                "pip install stable-baselines3"
+            ) from exc
+
+        assert self._algo_cls is not None
+        self.model = self._algo_cls(
+            policy,
+            env,
+            seed=seed,
+            verbose=0,
+            **sb3_kwargs,
+        )
+
+    def train(
+        self,
+        n_episodes: int = 500,
+        verbose: bool = True,
+        **_kwargs: object,
+    ) -> dict[str, Any]:
+        """Train for *n_episodes* via SB3's ``learn`` method.
+
+        SB3 operates in total timesteps, so we estimate the budget as
+        ``n_episodes × env.spec.max_episode_steps`` (or 500 fallback).
+        """
+        max_ep = getattr(self.env.spec, "max_episode_steps", None) or 500
+        total_timesteps = n_episodes * max_ep
+
+        # Use a callback to record per-episode statistics
+        self.model.learn(total_timesteps=total_timesteps, progress_bar=verbose)
+
+        # Evaluate to gather reward traces
+        self._evaluate(n_episodes, verbose)
+        return {
+            "episode_rewards": self.episode_rewards,
+            "episode_lengths": self.episode_lengths,
+            "total_steps": sum(self.episode_lengths),
+        }
+
+    def act(self, state: NDArray[np.floating]) -> NDArray[np.floating]:
+        action, _ = self.model.predict(state, deterministic=True)
+        return action
+
+    def _evaluate(self, n_episodes: int, verbose: bool) -> None:
+        pbar = tqdm(range(n_episodes), desc=self.__class__.__name__ + " eval", disable=not verbose)
+        for _ in pbar:
+            state, _ = self.env.reset()
+            ep_reward = 0.0
+            ep_len = 0
+            terminated = truncated = False
+            while not (terminated or truncated):
+                action = self.act(state)
+                state, reward, terminated, truncated, _ = self.env.step(action)
+                ep_reward += float(reward)
+                ep_len += 1
+            self.episode_rewards.append(ep_reward)
+            self.episode_lengths.append(ep_len)
+
+
+class PPOAgent(_SB3Agent):
+    """Proximal Policy Optimisation (PPO) via stable-baselines3."""
+
+    def __init__(self, env: gym.Env, seed: int | None = None, **kwargs: object) -> None:
+        from stable_baselines3 import PPO
+
+        self._algo_cls = PPO  # type: ignore[assignment]
+        super().__init__(env, seed=seed, **kwargs)
+
+
+class SACAgent(_SB3Agent):
+    """Soft Actor-Critic (SAC) via stable-baselines3."""
+
+    def __init__(self, env: gym.Env, seed: int | None = None, **kwargs: object) -> None:
+        from stable_baselines3 import SAC
+
+        self._algo_cls = SAC  # type: ignore[assignment]
+        super().__init__(env, seed=seed, **kwargs)