"""
core.py - Interfaz principal (Facade).
"""
from typing import List, Callable, Optional
import h5py
from .data import load_tripodpy_hdf5
from .composition import CompositionModel, SimpleWaterComposition, FunctionComposition
from .pebble_accretion import PebbleAccretionModule3
from .snowline import generate_rsnow_array
from .plotting import plot_hovmoller
from . import constants as c
[docs]
class PA3Py:
"""
Centraliza flujo de PA3Py: Datos HDF5, química, acreción y gráficas.
"""
def __init__(self, data_dir: str, comp_model: Optional[CompositionModel] = None):
"""
Inicializa la simulación y carga los datos.
Parámetros:
-----------
data_dir : str
Ruta a archivos HDF5.
comp_model : CompositionModel, opcional
Modelo de composición. Por defecto: Agua simple + migración snowline.
"""
self.data_dir = data_dir
self.disk = load_tripodpy_hdf5(data_dir)
self.comp = comp_model or SimpleWaterComposition(generate_rsnow_array(self.disk.times))
self.engine = PebbleAccretionModule3(self.disk, comp_model=self.comp)
[docs]
def set_custom_chemistry(self, user_func: Callable, species: Optional[List[str]] = None):
"""
Atajo para redefinir la química con una función Python.
Ejemplo:
sim.set_custom_chemistry(mi_funcion, ["silicatos", "H2O"])
"""
self.comp = FunctionComposition(user_func, species)
self.engine = PebbleAccretionModule3(self.disk, comp_model=self.comp)
[docs]
def run_growth(self, embryos_au: list, m_seed_me: float = 1e-3) -> dict:
"""
Corre la simulación de acreción para los embriones dados.
Parámetros:
-----------
embryos_au : list
Lista de radios iniciales en AU (ej: [1.0, 5.0, 15.0]).
m_seed_me : float
Masa semilla en Masas Terrestres. Default: 1e-3.
Retorna:
--------
dict
Resultados de la evolución de masa en el tiempo.
"""
results = self.engine.run_growth(embryos_au, M0_g=m_seed_me * c.M_EARTH)
self.engine.summary(results)
return results
[docs]
def plot_hovmoller(self, field: str = 'dust_Sigma', show_snowlines: bool = True, **kwargs):
"""
Genera el diagrama de Hovmöller (Radio vs Tiempo).
field puede ser 'dust_Sigma', 'gas_Sigma', o 'epsilon'.
"""
return plot_hovmoller(self.disk, field=field, show_snowlines=show_snowlines, **kwargs)
[docs]
def calculate_isolation_mass_map(self):
"""Calcula el mapa teórico de masa de aislamiento en todo el disco."""
return self.engine.calculate_isolation_mass_map()
[docs]
def plot_population(self, results: dict, **kwargs):
"""
Grafica la población sintética de planetas (Masa Final vs Posición Inicial).
"""
from .plotting import plot_population
M_iso_map = self.calculate_isolation_mass_map()
return plot_population(self.disk, results, M_iso_map=M_iso_map, **kwargs)
[docs]
def save_results(self, results: dict, filename: str):
"""Guarda tracks de crecimiento en HDF5."""
with h5py.File(filename, 'w') as f:
# Encode species como bytes para compatibilidad h5py
f.attrs['tracked_species'] = [s.encode('ascii') for s in self.engine.tracked_species]
f.attrs['M_EARTH'] = c.M_EARTH
tracks = f.create_group('tracks')
for r_au, track_data in results.items():
dset = tracks.create_dataset(f"embryo_{r_au:.4f}_AU", data=track_data)
dset.attrs['r_au'] = float(r_au)
[docs]
@staticmethod
def load_results(filename: str) -> tuple[dict, list]:
"""Carga tracks de crecimiento desde HDF5. Retorna (results_dict, species_list)."""
results = {}
with h5py.File(filename, 'r') as f:
raw_species = list(f.attrs['tracked_species'])
tracked_species = [s.decode('ascii') if isinstance(s, bytes) else s for s in raw_species]
tracks = f['tracks']
for key in tracks.keys():
r_au = float(tracks[key].attrs['r_au'])
results[r_au] = tracks[key][:]
return results, tracked_species