Update galpy/mcp_output/mcp_plugin/mcp_service.py

galpy/mcp_output/mcp_plugin/mcp_service.py

@@ -1,82 +1,415 @@
 import os
 import sys
 
-# Add the local source directory to sys.path
-source_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "source")
-if source_path not in sys.path:
-    sys.path.insert(0, source_path)
+# NOTE: galpy should be installed via pip, no need to add source path
+# This allows using the compiled Cython extensions properly
 
 from fastmcp import FastMCP
+import numpy as np
+
+# Import core modules from galpy
 from galpy.orbit import Orbit
-from galpy.potential import Potential
-from galpy.actionAngle import actionAngleIsochrone
+from galpy.potential import (
+    MiyamotoNagaiPotential,
+    NFWPotential,
+    LogarithmicHaloPotential,
+    HernquistPotential,
+    PlummerPotential,
+    IsochronePotential,
+    KeplerPotential,
+    evaluatePotentials,
+    vcirc,
+    epifreq,
+    verticalfreq,
+    omegac,
+)
+from galpy.actionAngle import (
+    actionAngleIsochrone,
+    actionAngleStaeckel,
+    actionAngleSpherical,
+    estimateDeltaStaeckel,
+)
 
 # Create the FastMCP service application
 mcp = FastMCP("galpy_service")
 
-@mcp.tool(name="calculate_orbit", description="Calculate the orbit of a celestial object.")
-def calculate_orbit(R: float, vR: float, vT: float, z: float, vz: float, phi: float) -> dict:
+# ============= Potential Creation Tools =============
+
+def _create_potential(potential_type: str, **params):
+    """Helper function to create potential objects."""
+    potential_map = {
+        "miyamoto_nagai": MiyamotoNagaiPotential,
+        "nfw": NFWPotential,
+        "logarithmic": LogarithmicHaloPotential,
+        "hernquist": HernquistPotential,
+        "plummer": PlummerPotential,
+        "isochrone": IsochronePotential,
+        "kepler": KeplerPotential,
+    }
+    if potential_type.lower() not in potential_map:
+        raise ValueError(f"Unknown potential type: {potential_type}. Available: {list(potential_map.keys())}")
+    return potential_map[potential_type.lower()](**params)
+
+@mcp.tool(name="create_potential", description="Create a gravitational potential model. Supported types: miyamoto_nagai, nfw, logarithmic, hernquist, plummer, isochrone, kepler.")
+def create_potential(potential_type: str, amp: float = 1.0, a: float = 1.0, b: float = 0.1) -> dict:
+    """
+    Create a gravitational potential model.
+
+    Parameters:
+    - potential_type: str - Type of potential (miyamoto_nagai, nfw, logarithmic, hernquist, plummer, isochrone, kepler)
+    - amp: float - Amplitude/mass parameter (default: 1.0)
+    - a: float - Scale length parameter (default: 1.0)
+    - b: float - Scale height parameter for disk potentials (default: 0.1)
+
+    Returns:
+    - dict: Contains success and potential info or error
+    """
+    try:
+        pot = _create_potential(potential_type, amp=amp, a=a, b=b) if potential_type.lower() == "miyamoto_nagai" else _create_potential(potential_type, amp=amp, a=a)
+        return {"success": True, "result": f"Created {potential_type} potential with amp={amp}, a={a}"}
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+@mcp.tool(name="evaluate_potential", description="Evaluate the gravitational potential at a given position (R, z).")
+def evaluate_potential(R: float, z: float, potential_type: str = "miyamoto_nagai", amp: float = 1.0, a: float = 1.0, b: float = 0.1) -> dict:
+    """
+    Evaluate the gravitational potential at a given point.
+
+    Parameters:
+    - R: float - Cylindrical Galactocentric radius (in natural units, 8 kpc = 1)
+    - z: float - Vertical height above the plane (in natural units)
+    - potential_type: str - Type of potential to use
+    - amp: float - Amplitude/mass parameter
+    - a: float - Scale length
+    - b: float - Scale height (for disk potentials)
+
+    Returns:
+    - dict: Contains success and potential value or error
+    """
+    try:
+        if potential_type.lower() == "miyamoto_nagai":
+            pot = MiyamotoNagaiPotential(amp=amp, a=a, b=b)
+        elif potential_type.lower() == "nfw":
+            pot = NFWPotential(amp=amp, a=a)
+        elif potential_type.lower() == "logarithmic":
+            pot = LogarithmicHaloPotential(amp=amp)
+        else:
+            pot = MiyamotoNagaiPotential(amp=amp, a=a, b=b)
+        
+        result = float(pot(R, z))
+        return {"success": True, "result": result, "R": R, "z": z, "potential_type": potential_type}
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+# ============= Orbit Tools =============
+
+@mcp.tool(name="calculate_orbit", description="Calculate and integrate the orbit of a celestial object in a gravitational potential.")
+def calculate_orbit(R: float, vR: float, vT: float, z: float, vz: float, phi: float, 
+                   t_end: float = 10.0, n_steps: int = 1000,
+                   potential_type: str = "miyamoto_nagai") -> dict:
     """
     Calculate the orbit of a celestial object given initial conditions.
 
     Parameters:
-    - R: float - Radial distance
-    - vR: float - Radial velocity
-    - vT: float - Tangential velocity
-    - z: float - Vertical position
-    - vz: float - Vertical velocity
-    - phi: float - Azimuthal angle
+    - R: float - Initial radial distance (in natural units, 8 kpc = 1)
+    - vR: float - Initial radial velocity (in natural units, 220 km/s = 1)
+    - vT: float - Initial tangential velocity
+    - z: float - Initial vertical position
+    - vz: float - Initial vertical velocity
+    - phi: float - Initial azimuthal angle (radians)
+    - t_end: float - End time for integration (in natural units, ~4.7 Gyr = 1)
+    - n_steps: int - Number of time steps
+    - potential_type: str - Type of gravitational potential
 
     Returns:
-    - dict: Contains success, result, or error fields
+    - dict: Contains success, orbit data (R, z, phi arrays), or error
     """
     try:
+        # Create potential
+        if potential_type.lower() == "miyamoto_nagai":
+            pot = MiyamotoNagaiPotential(amp=1.0, a=1.0, b=0.1)
+        elif potential_type.lower() == "nfw":
+            pot = NFWPotential(amp=1.0, a=1.0)
+        elif potential_type.lower() == "logarithmic":
+            pot = LogarithmicHaloPotential(amp=1.0)
+        else:
+            pot = MiyamotoNagaiPotential(amp=1.0, a=1.0, b=0.1)
+        
+        # Create orbit with initial conditions
         orbit = Orbit([R, vR, vT, z, vz, phi])
-        result = orbit.integrate()
+        
+        # Time array
+        ts = np.linspace(0, t_end, n_steps)
+        
+        # Integrate
+        orbit.integrate(ts, pot)
+        
+        # Extract orbit data (sample every 10th point to reduce output size)
+        sample_rate = max(1, n_steps // 100)
+        result = {
+            "R": orbit.R(ts[::sample_rate]).tolist(),
+            "z": orbit.z(ts[::sample_rate]).tolist(),
+            "phi": orbit.phi(ts[::sample_rate]).tolist(),
+            "t": ts[::sample_rate].tolist(),
+        }
         return {"success": True, "result": result}
     except Exception as e:
         return {"success": False, "error": str(e)}
 
-@mcp.tool(name="evaluate_potential", description="Evaluate the gravitational potential at a point.")
-def evaluate_potential(R: float, z: float) -> dict:
+@mcp.tool(name="calculate_orbit_energy", description="Calculate the energy of an orbit at given initial conditions.")
+def calculate_orbit_energy(R: float, vR: float, vT: float, z: float, vz: float, phi: float,
+                          potential_type: str = "miyamoto_nagai") -> dict:
     """
-    Evaluate the gravitational potential at a given point.
+    Calculate the total energy of an orbit.
 
     Parameters:
-    - R: float - Radial distance
-    - z: float - Vertical position
+    - R, vR, vT, z, vz, phi: Initial conditions for the orbit
+    - potential_type: str - Type of gravitational potential
 
     Returns:
-    - dict: Contains success, result, or error fields
+    - dict: Contains success and energy value or error
     """
     try:
-        potential = Potential()
-        result = potential(R, z)
-        return {"success": True, "result": result}
+        if potential_type.lower() == "miyamoto_nagai":
+            pot = MiyamotoNagaiPotential(amp=1.0, a=1.0, b=0.1)
+        elif potential_type.lower() == "nfw":
+            pot = NFWPotential(amp=1.0, a=1.0)
+        else:
+            pot = MiyamotoNagaiPotential(amp=1.0, a=1.0, b=0.1)
+        
+        orbit = Orbit([R, vR, vT, z, vz, phi])
+        orbit.integrate([0, 0.01], pot)  # Need to integrate briefly to access energy
+        energy = float(orbit.E(pot=pot))
+        
+        return {"success": True, "result": energy, "unit": "natural units (220 km/s)^2"}
     except Exception as e:
         return {"success": False, "error": str(e)}
 
-@mcp.tool(name="compute_action_angle", description="Compute action-angle variables for a given orbit.")
-def compute_action_angle(R: float, vR: float, vT: float, z: float, vz: float, phi: float) -> dict:
+@mcp.tool(name="calculate_angular_momentum", description="Calculate the angular momentum of an orbit.")
+def calculate_angular_momentum(R: float, vR: float, vT: float, z: float, vz: float, phi: float) -> dict:
+    """
+    Calculate the angular momentum vector of an orbit.
+
+    Parameters:
+    - R, vR, vT, z, vz, phi: Initial conditions for the orbit
+
+    Returns:
+    - dict: Contains success and angular momentum (Lx, Ly, Lz) or error
+    """
+    try:
+        orbit = Orbit([R, vR, vT, z, vz, phi])
+        L = orbit.L()
+        
+        return {
+            "success": True, 
+            "result": {
+                "Lx": float(L[0]),
+                "Ly": float(L[1]),
+                "Lz": float(L[2])
+            },
+            "unit": "natural units (8 kpc * 220 km/s)"
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+@mcp.tool(name="calculate_orbital_parameters", description="Calculate orbital parameters: eccentricity, apocenter, pericenter.")
+def calculate_orbital_parameters(R: float, vR: float, vT: float, z: float, vz: float, phi: float,
+                                 potential_type: str = "miyamoto_nagai") -> dict:
+    """
+    Calculate orbital parameters from initial conditions.
+
+    Parameters:
+    - R, vR, vT, z, vz, phi: Initial conditions for the orbit
+    - potential_type: str - Type of gravitational potential
+
+    Returns:
+    - dict: Contains eccentricity, apocenter radius, pericenter radius, or error
+    """
+    try:
+        if potential_type.lower() == "miyamoto_nagai":
+            pot = MiyamotoNagaiPotential(amp=1.0, a=1.0, b=0.1)
+        else:
+            pot = MiyamotoNagaiPotential(amp=1.0, a=1.0, b=0.1)
+        
+        orbit = Orbit([R, vR, vT, z, vz, phi])
+        ts = np.linspace(0, 20, 2000)  # Integrate for longer to get good orbital params
+        orbit.integrate(ts, pot)
+        
+        ecc = float(orbit.e())
+        rap = float(orbit.rap())
+        rperi = float(orbit.rperi())
+        
+        return {
+            "success": True,
+            "result": {
+                "eccentricity": ecc,
+                "apocenter": rap,
+                "pericenter": rperi
+            },
+            "unit": "natural units (8 kpc for distances)"
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+# ============= Action-Angle Tools =============
+
+@mcp.tool(name="compute_action_angle", description="Compute action-angle variables for a given orbit using the isochrone approximation.")
+def compute_action_angle(R: float, vR: float, vT: float, z: float, vz: float, phi: float,
+                        method: str = "isochrone") -> dict:
     """
     Compute action-angle variables for a given orbit.
 
     Parameters:
-    - R: float - Radial distance
-    - vR: float - Radial velocity
-    - vT: float - Tangential velocity
-    - z: float - Vertical position
-    - vz: float - Vertical velocity
-    - phi: float - Azimuthal angle
+    - R, vR, vT, z, vz, phi: Orbital phase-space coordinates
+    - method: str - Method to use: 'isochrone', 'staeckel', or 'spherical'
 
     Returns:
-    - dict: Contains success, result, or error fields
+    - dict: Contains actions (JR, Lz, Jz) and angles or error
     """
     try:
         orbit = Orbit([R, vR, vT, z, vz, phi])
-        action_angle = actionAngleIsochrone()
-        result = action_angle(orbit)
-        return {"success": True, "result": result}
+        
+        if method.lower() == "isochrone":
+            aA = actionAngleIsochrone(ip=IsochronePotential(amp=1.0, b=0.9))
+        elif method.lower() == "staeckel":
+            pot = MiyamotoNagaiPotential(amp=1.0, a=1.0, b=0.1)
+            delta = estimateDeltaStaeckel(pot, R, z)
+            aA = actionAngleStaeckel(pot=pot, delta=delta)
+        elif method.lower() == "spherical":
+            pot = LogarithmicHaloPotential(amp=1.0)
+            aA = actionAngleSpherical(pot=pot)
+        else:
+            aA = actionAngleIsochrone(ip=IsochronePotential(amp=1.0, b=0.9))
+        
+        # Compute actions
+        jr, lz, jz = aA(R, vR, vT, z, vz)
+        
+        return {
+            "success": True,
+            "result": {
+                "JR": float(jr),
+                "Lz": float(lz),
+                "Jz": float(jz)
+            },
+            "method": method,
+            "unit": "natural units (8 kpc * 220 km/s)"
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+# ============= Galaxy Dynamics Tools =============
+
+@mcp.tool(name="calculate_circular_velocity", description="Calculate the circular velocity at a given radius in a potential.")
+def calculate_circular_velocity(R: float, potential_type: str = "miyamoto_nagai", 
+                                amp: float = 1.0, a: float = 1.0, b: float = 0.1) -> dict:
+    """
+    Calculate circular velocity at radius R.
+
+    Parameters:
+    - R: float - Galactocentric radius
+    - potential_type: str - Type of potential
+    - amp, a, b: Potential parameters
+
+    Returns:
+    - dict: Contains circular velocity or error
+    """
+    try:
+        if potential_type.lower() == "miyamoto_nagai":
+            pot = MiyamotoNagaiPotential(amp=amp, a=a, b=b)
+        elif potential_type.lower() == "nfw":
+            pot = NFWPotential(amp=amp, a=a)
+        elif potential_type.lower() == "logarithmic":
+            pot = LogarithmicHaloPotential(amp=amp)
+        else:
+            pot = MiyamotoNagaiPotential(amp=amp, a=a, b=b)
+        
+        vc = float(vcirc(pot, R))
+        
+        return {
+            "success": True,
+            "result": vc,
+            "R": R,
+            "unit": "natural units (220 km/s = 1)"
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+@mcp.tool(name="calculate_frequencies", description="Calculate orbital frequencies (epicyclic, vertical, circular) at a given radius.")
+def calculate_frequencies(R: float, potential_type: str = "miyamoto_nagai",
+                         amp: float = 1.0, a: float = 1.0, b: float = 0.1) -> dict:
+    """
+    Calculate orbital frequencies at radius R.
+
+    Parameters:
+    - R: float - Galactocentric radius
+    - potential_type: str - Type of potential
+    - amp, a, b: Potential parameters
+
+    Returns:
+    - dict: Contains epicyclic, vertical, and circular frequencies or error
+    """
+    try:
+        if potential_type.lower() == "miyamoto_nagai":
+            pot = MiyamotoNagaiPotential(amp=amp, a=a, b=b)
+        elif potential_type.lower() == "nfw":
+            pot = NFWPotential(amp=amp, a=a)
+        else:
+            pot = MiyamotoNagaiPotential(amp=amp, a=a, b=b)
+        
+        kappa = float(epifreq(pot, R))      # Epicyclic frequency
+        nu = float(verticalfreq(pot, R))    # Vertical frequency  
+        omega = float(omegac(pot, R))       # Circular frequency
+        
+        return {
+            "success": True,
+            "result": {
+                "epicyclic_frequency": kappa,
+                "vertical_frequency": nu,
+                "circular_frequency": omega
+            },
+            "R": R,
+            "unit": "natural units (1/time, where time unit ~ 4.7 Gyr)"
+        }
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+@mcp.tool(name="calculate_rotation_curve", description="Calculate the rotation curve of a galaxy model.")
+def calculate_rotation_curve(R_min: float = 0.1, R_max: float = 3.0, n_points: int = 50,
+                            potential_type: str = "miyamoto_nagai",
+                            amp: float = 1.0, a: float = 1.0, b: float = 0.1) -> dict:
+    """
+    Calculate the rotation curve (circular velocity vs radius).
+
+    Parameters:
+    - R_min, R_max: float - Range of radii to calculate
+    - n_points: int - Number of points
+    - potential_type: str - Type of potential
+    - amp, a, b: Potential parameters
+
+    Returns:
+    - dict: Contains arrays of R and v_circ or error
+    """
+    try:
+        if potential_type.lower() == "miyamoto_nagai":
+            pot = MiyamotoNagaiPotential(amp=amp, a=a, b=b)
+        elif potential_type.lower() == "nfw":
+            pot = NFWPotential(amp=amp, a=a)
+        elif potential_type.lower() == "logarithmic":
+            pot = LogarithmicHaloPotential(amp=amp)
+        else:
+            pot = MiyamotoNagaiPotential(amp=amp, a=a, b=b)
+        
+        Rs = np.linspace(R_min, R_max, n_points)
+        vcs = [float(vcirc(pot, r)) for r in Rs]
+        
+        return {
+            "success": True,
+            "result": {
+                "R": Rs.tolist(),
+                "v_circular": vcs
+            },
+            "unit": "R in 8 kpc, v in 220 km/s"
+        }
     except Exception as e:
         return {"success": False, "error": str(e)}