From b161820b7bf5f97d0a39cf54eb81d608c22b4889 Mon Sep 17 00:00:00 2001 From: Silen Naihin Date: Thu, 18 Dec 2025 15:20:20 -0500 Subject: [PATCH 1/3] pydantic --- .../lux/projects/alab/schemas/__init__.py | 36 +++ .../lux/projects/alab/schemas/base.py | 69 ++++++ .../alab/schemas/experiment_elements.py | 30 +++ .../lux/projects/alab/schemas/experiments.py | 227 ++++++++++++++++++ .../lux/projects/alab/schemas/powder_doses.py | 76 ++++++ .../projects/alab/schemas/temperature_logs.py | 51 ++++ .../projects/alab/schemas/workflow_tasks.py | 55 +++++ .../projects/alab/schemas/xrd_data_points.py | 38 +++ .../lux/projects/alab/schemas/xrd_phases.py | 64 +++++ .../projects/alab/schemas/xrd_refinements.py | 89 +++++++ 10 files changed, 735 insertions(+) create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/__init__.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/base.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/experiment_elements.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/experiments.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/powder_doses.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/temperature_logs.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/workflow_tasks.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/xrd_data_points.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/xrd_phases.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/xrd_refinements.py diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/__init__.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/__init__.py new file mode 100644 index 000000000..e30cfc64a --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/__init__.py @@ -0,0 +1,36 @@ +""" +A-Lab Pydantic Schemas + +This package contains Pydantic schemas for all A-Lab parquet tables. +These schemas are the source of truth for data validation. + +Each schema corresponds to one parquet file. +Integrates team's validation patterns (constraints, Literal types) from results_schema.py. +""" + +from .base import ExcludeFromUpload +from .experiments import Experiment +from .experiment_elements import ExperimentElement +from .powder_doses import PowderDose, PowderItem +from .temperature_logs import TemperatureLog, TemperatureLogEntry +from .workflow_tasks import WorkflowTask +from .xrd_data_points import XRDDataPoint +from .xrd_refinements import XRDRefinement +from .xrd_phases import XRDPhase + +__all__ = [ + # Base + "ExcludeFromUpload", + # Parquet table schemas (one per .parquet file) + "Experiment", + "ExperimentElement", + "PowderDose", + "PowderItem", + "TemperatureLogEntry", + "TemperatureLog", + "WorkflowTask", + "XRDDataPoint", + "XRDRefinement", + "XRDPhase", +] + diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/base.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/base.py new file mode 100644 index 000000000..5bae46d12 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/base.py @@ -0,0 +1,69 @@ +""" +Base utilities for A-Lab Pydantic schemas. + +Provides common types, validators, and field utilities. +""" + +from typing import Any, Dict +from pydantic import Field + + +def ExcludeFromUpload( + default: Any = None, + description: str = "", + **kwargs +) -> Any: + """ + Field that should NOT be uploaded to MPContribs. + + Use this for sensitive data that must remain private until publication. + Examples: weight_collected, mass measurements that are embargoed. + + Usage: + weight_collected: float | None = ExcludeFromUpload( + description="Weight of powder collected (embargoed)" + ) + """ + return Field( + default=default, + description=description, + json_schema_extra={"exclude_from_upload": True}, + **kwargs + ) + + +def get_uploadable_fields(model_class) -> list[str]: + """ + Get list of fields that should be uploaded to MPContribs. + + Args: + model_class: Pydantic model class + + Returns: + List of field names that are NOT marked with exclude_from_upload + """ + uploadable = [] + for field_name, field_info in model_class.model_fields.items(): + extra = field_info.json_schema_extra or {} + if not extra.get("exclude_from_upload", False): + uploadable.append(field_name) + return uploadable + + +def get_excluded_fields(model_class) -> list[str]: + """ + Get list of fields that should NOT be uploaded to MPContribs. + + Args: + model_class: Pydantic model class + + Returns: + List of field names that ARE marked with exclude_from_upload + """ + excluded = [] + for field_name, field_info in model_class.model_fields.items(): + extra = field_info.json_schema_extra or {} + if extra.get("exclude_from_upload", False): + excluded.append(field_name) + return excluded + diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/experiment_elements.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/experiment_elements.py new file mode 100644 index 000000000..ad90a56dd --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/experiment_elements.py @@ -0,0 +1,30 @@ +""" +Experiment Elements Schema + +Elements present in each experiment (1:N relationship). +Maps to: experiment_elements.parquet +""" + +from pydantic import BaseModel, Field + + +class ExperimentElement(BaseModel, extra="forbid"): + """ + Element present in an experiment. + + Each experiment can have multiple elements (1:N relationship). + """ + + experiment_id: str = Field( + description="Reference to parent experiment" + ) + + element_symbol: str = Field( + description="Element symbol (e.g., Na, Mg, O)" + ) + + target_atomic_percent: float | None = Field( + default=None, + description="Target atomic percentage of this element" + ) + diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/experiments.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/experiments.py new file mode 100644 index 000000000..5ca1ffdd3 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/experiments.py @@ -0,0 +1,227 @@ +""" +Experiment Schema (Consolidated) + +This is the main experiment table with ALL 1:1 data merged. +Contains ~45 columns from: experiments + heating + recovery + xrd + finalization + dosing. + +Maps to: experiments.parquet +""" + +from datetime import datetime +from typing import Literal +from pydantic import BaseModel, Field + +# Import ExcludeFromUpload utility +try: + from .base import ExcludeFromUpload +except ImportError: + # When imported dynamically, use absolute import + import sys + from pathlib import Path + sys.path.insert(0, str(Path(__file__).parent)) + from base import ExcludeFromUpload + + +class Experiment(BaseModel, extra="forbid"): + """ + Main experiment schema with all 1:1 data consolidated. + + This is the primary table - one row per experiment. + All related 1:1 data (heating, recovery, xrd, finalization, dosing) is merged here. + """ + + # === Core experiment fields === + experiment_id: str = Field( + description="Unique experiment identifier (MongoDB _id)" + ) + + name: str = Field( + description="Experiment name (e.g., NSC_249, MINES_12)" + ) + + experiment_type: str = Field( + description="Root experiment type (NSC, Na, PG, MINES, TRI)" + ) + + experiment_subgroup: str | None = Field( + default=None, + description="Experiment subgroup (e.g., NSC_249, Na_123)" + ) + + target_formula: str = Field( + description="Target chemical formula" + ) + + last_updated: datetime = Field( + description="Last modification timestamp" + ) + + status: Literal["completed", "error", "active", "unknown"] = Field( + description="Workflow status" + ) + + notes: str | None = Field( + default=None, + description="Optional notes about the experiment" + ) + + # === Heating fields (prefix: heating_) === + heating_method: Literal["standard", "atmosphere", "manual", "none"] | None = Field( + default=None, + description="Heating method used" + ) + + heating_temperature: float | None = Field( + default=None, + description="Target heating temperature in °C" + ) + + heating_time: float | None = Field( + default=None, + description="Heating duration in minutes" + ) + + heating_cooling_rate: float | None = Field( + default=None, + description="Cooling rate in °C/min" + ) + + heating_atmosphere: str | None = Field( + default=None, + description="Atmosphere used during heating (e.g., N2, Ar, Air)" + ) + + heating_flow_rate_ml_min: float | None = Field( + default=None, + description="Gas flow rate during heating in mL/min" + ) + + heating_low_temp_calcination: bool | None = Field( + default=None, + description="Whether low temperature calcination was used" + ) + + # === Recovery fields (prefix: recovery_) === + recovery_total_dosed_mass_mg: float | None = Field( + default=None, + description="Total mass of all powders dosed in mg" + ) + + # EXCLUDED FROM UPLOAD per team request + recovery_weight_collected_mg: float | None = ExcludeFromUpload( + description="Weight of powder collected after heating in mg (EMBARGOED)" + ) + + recovery_yield_percent: float | None = Field( + default=None, + description="Recovery yield (collected / dosed * 100)" + ) + + recovery_initial_crucible_weight_mg: float | None = Field( + default=None, + description="Initial crucible weight before experiment in mg" + ) + + recovery_failure_classification: str | None = Field( + default=None, + description="Classification of any failure during recovery" + ) + + # === XRD measurement fields (prefix: xrd_) === + xrd_sampleid_in_aeris: str | None = Field( + default=None, + description="Sample ID in Aeris XRD system" + ) + + xrd_holder_index: int | None = Field( + default=None, + description="XRD sample holder position index" + ) + + # EXCLUDED FROM UPLOAD per team request + xrd_total_mass_dispensed_mg: float | None = ExcludeFromUpload( + description="Mass dispensed for XRD measurement in mg (EMBARGOED)" + ) + + xrd_met_target_mass: bool | None = Field( + default=None, + description="Whether target mass was achieved for XRD" + ) + + # === Finalization fields (prefix: finalization_) === + finalization_decoded_sample_id: str | None = Field( + default=None, + description="Decoded sample ID from barcode" + ) + + finalization_successful_labeling: bool | None = Field( + default=None, + description="Whether sample was successfully labeled" + ) + + finalization_storage_location: str | None = Field( + default=None, + description="Final storage location of sample" + ) + + # === Dosing fields (prefix: dosing_) === + dosing_crucible_position: int | None = Field( + default=None, + description="Crucible position in rack", + ge=1, + le=4 + ) + + dosing_crucible_sub_rack: Literal["SubRackA", "SubRackB", "SubRackC", "SubRackD"] | None = Field( + default=None, + description="Sub-rack identifier" + ) + + dosing_mixing_pot_position: int | None = Field( + default=None, + description="Mixing pot position", + ge=1, + le=16 + ) + + dosing_ethanol_dispense_volume: int | None = Field( + default=None, + description="Volume of ethanol dispensed in µL", + ge=0 + ) + + dosing_target_transfer_volume: int | None = Field( + default=None, + description="Target transfer volume in µL", + ge=0 + ) + + dosing_actual_transfer_mass: float | None = Field( + default=None, + description="Actual mass transferred in g", + ge=0 + ) + + dosing_dac_duration: int | None = Field( + default=None, + description="DAC duration in seconds", + ge=0 + ) + + dosing_dac_speed: int | None = Field( + default=None, + description="DAC rotation speed in rpm", + ge=0 + ) + + dosing_actual_heat_duration: int | None = Field( + default=None, + description="Actual heating duration during dosing in seconds", + ge=0 + ) + + dosing_end_reason: str | None = Field( + default=None, + description="Reason for ending dosing session" + ) + diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/powder_doses.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/powder_doses.py new file mode 100644 index 000000000..f2f376f32 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/powder_doses.py @@ -0,0 +1,76 @@ +""" +Powder Doses Schema + +Individual powder doses for each experiment (1:N relationship). +Maps to: powder_doses.parquet + +NOTE: This is distinct from PowderDosingSampleResult which is the raw +MongoDB structure. This is the flattened parquet schema. +""" + +from pydantic import BaseModel, Field + + +class PowderItem(BaseModel, extra="forbid"): + """Schema for a powder item within a dose.""" + + PowderName: str | None = Field( + default=None, + description="Name of the powder" + ) + TargetMass: float | None = Field( + default=None, + description="Target mass in mg", + ge=0 + ) + Doses: list[dict] | None = Field( + default=None, + description="List of individual doses" + ) + + +class PowderDose(BaseModel, extra="forbid"): + """ + Individual powder dose event. + + Each experiment can have multiple powder doses (1:N relationship). + This is the flattened representation from the nested Powders[].Doses[] structure. + """ + + experiment_id: str = Field( + description="Reference to parent experiment" + ) + + powder_name: str = Field( + description="Name of the powder material" + ) + + target_mass: float = Field( + description="Target mass in mg", + ge=0 + ) + + actual_mass: float = Field( + description="Actual mass dispensed in mg", + ge=0 + ) + + accuracy_percent: float = Field( + description="Dosing accuracy (actual/target * 100)" + ) + + dose_sequence: int = Field( + description="Sequence number within the experiment", + ge=0 + ) + + head_position: int | None = Field( + default=None, + description="Dispenser head position" + ) + + dose_timestamp: str | None = Field( + default=None, + description="Timestamp of the dose event" + ) + diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/temperature_logs.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/temperature_logs.py new file mode 100644 index 000000000..980c7af9d --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/temperature_logs.py @@ -0,0 +1,51 @@ +""" +Temperature Logs Schema + +Temperature readings during heating (1:N relationship). +Maps to: temperature_logs.parquet + +NOTE: This is the flattened parquet schema. The raw MongoDB schema uses +nested arrays (time_minutes[], temperature_celsius[]). +""" + +from pydantic import BaseModel, Field + + +class TemperatureLogEntry(BaseModel, extra="forbid"): + """ + Single temperature reading. + + Each experiment can have thousands of temperature log entries (1:N relationship). + """ + + experiment_id: str = Field( + description="Reference to parent experiment" + ) + + sequence_number: int = Field( + description="Sequence number in the time series", + ge=0 + ) + + time_minutes: float = Field( + description="Time elapsed since heating start in minutes" + ) + + temperature_celsius: float = Field( + description="Temperature reading in °C" + ) + + +# For backward compatibility with nested structure +class TemperatureLog(BaseModel, extra="forbid"): + """Temperature log with array structure (MongoDB format).""" + + time_minutes: list[float] | None = Field( + default=None, + description="Time elapsed in minutes since the start of the heating process" + ) + temperature_celsius: list[float] | None = Field( + default=None, + description="Temperature of the sample in Celsius" + ) + diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/workflow_tasks.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/workflow_tasks.py new file mode 100644 index 000000000..7b6c6a1d2 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/workflow_tasks.py @@ -0,0 +1,55 @@ +""" +Workflow Tasks Schema + +Task execution history for each experiment (1:N relationship). +Maps to: workflow_tasks.parquet +""" + +from datetime import datetime +from typing import Literal +from pydantic import BaseModel, Field + + +class WorkflowTask(BaseModel, extra="forbid"): + """ + Single task in the experiment workflow. + + Each experiment has multiple tasks (1:N relationship). + Tasks include: PowderDosing, Heating, RecoverPowder, Diffraction, Ending, etc. + """ + + experiment_id: str = Field( + description="Reference to parent experiment" + ) + + task_id: str = Field( + description="Unique task identifier" + ) + + task_type: str = Field( + description="Type of task (PowderDosing, Heating, RecoverPowder, Diffraction, Ending)" + ) + + status: Literal["completed", "error", "pending", "running"] = Field( + description="Task execution status" + ) + + created_at: datetime = Field( + description="When the task was created" + ) + + started_at: datetime | None = Field( + default=None, + description="When the task started execution" + ) + + completed_at: datetime | None = Field( + default=None, + description="When the task completed" + ) + + message: str | None = Field( + default=None, + description="Task result message or error" + ) + diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/xrd_data_points.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/xrd_data_points.py new file mode 100644 index 000000000..4db376e72 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/xrd_data_points.py @@ -0,0 +1,38 @@ +""" +XRD Data Points Schema + +Raw XRD diffraction pattern data (1:N relationship). +Maps to: xrd_data_points.parquet + +NOTE: This is the flattened parquet schema. The raw MongoDB schema uses +nested arrays (twotheta[], counts[]). +""" + +from pydantic import BaseModel, Field + + +class XRDDataPoint(BaseModel, extra="forbid"): + """ + Single XRD data point (2θ, counts). + + Each experiment can have ~8000 data points (1:N relationship). + This is the flattened representation from twotheta[] and counts[] arrays. + """ + + experiment_id: str = Field( + description="Reference to parent experiment" + ) + + point_index: int = Field( + description="Index in the diffraction pattern", + ge=0 + ) + + twotheta: float = Field( + description="2θ angle in degrees" + ) + + counts: float = Field( + description="Intensity counts at this angle" + ) + diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/xrd_phases.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/xrd_phases.py new file mode 100644 index 000000000..dccd7ef92 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/xrd_phases.py @@ -0,0 +1,64 @@ +""" +XRD Phases Schema + +Identified crystal phases from DARA analysis (1:N relationship with experiments). +Maps to: xrd_phases.parquet +""" + +from pydantic import BaseModel, Field + + +class XRDPhase(BaseModel, extra="forbid"): + """ + Identified crystal phase from XRD refinement. + + Each experiment can have multiple phases (1:N relationship). + """ + + experiment_id: str = Field( + description="Reference to parent experiment" + ) + + experiment_name: str = Field( + description="Experiment name (for convenience)" + ) + + phase_name: str = Field( + description="Name of the identified phase" + ) + + spacegroup: str = Field( + description="Space group of the phase (e.g., Fm-3m)" + ) + + weight_fraction: float = Field( + description="Weight fraction of this phase (0-1)", + ge=0, + le=1 + ) + + weight_fraction_error: float | None = Field( + default=None, + description="Error in weight fraction" + ) + + lattice_a_nm: float | None = Field( + default=None, + description="Lattice parameter a in nm" + ) + + lattice_b_nm: float | None = Field( + default=None, + description="Lattice parameter b in nm" + ) + + lattice_c_nm: float | None = Field( + default=None, + description="Lattice parameter c in nm" + ) + + r_phase: float | None = Field( + default=None, + description="R-factor for this phase" + ) + diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/xrd_refinements.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/xrd_refinements.py new file mode 100644 index 000000000..9087b1136 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/schemas/xrd_refinements.py @@ -0,0 +1,89 @@ +""" +XRD Refinements Schema + +DARA Rietveld refinement analysis results (1:1 relationship with experiments). +Maps to: xrd_refinements.parquet +""" + +from pydantic import BaseModel, Field + + +class XRDRefinement(BaseModel, extra="forbid"): + """ + XRD Rietveld refinement result from DARA analysis. + + One result per experiment (1:1 relationship). + """ + + experiment_id: str = Field( + description="Reference to parent experiment" + ) + + experiment_name: str = Field( + description="Experiment name (for convenience)" + ) + + success: bool = Field( + description="Whether the refinement was successful" + ) + + error: str | None = Field( + default=None, + description="Error message if refinement failed" + ) + + error_type: str | None = Field( + default=None, + description="Classification of error type" + ) + + rwp: float | None = Field( + default=None, + description="Weighted profile R-factor (goodness of fit)" + ) + + rp: float | None = Field( + default=None, + description="Profile R-factor" + ) + + rexp: float | None = Field( + default=None, + description="Expected R-factor" + ) + + num_phases: int = Field( + default=0, + description="Number of phases identified" + ) + + chemical_system: str | None = Field( + default=None, + description="Chemical system (e.g., Na-Mg-P-O)" + ) + + target_formula: str | None = Field( + default=None, + description="Target formula" + ) + + analysis_timestamp: str | None = Field( + default=None, + description="When the analysis was performed" + ) + + mode: str | None = Field( + default=None, + description="Analysis mode" + ) + + wmin: int | None = Field( + default=None, + description="Minimum 2θ angle for refinement" + ) + + wmax: int | None = Field( + default=None, + description="Maximum 2θ angle for refinement" + ) + From 089ede3b61c2a2853c451b795abb2d9c6b7e9850 Mon Sep 17 00:00:00 2001 From: Silen Naihin Date: Thu, 18 Dec 2025 15:25:03 -0500 Subject: [PATCH 2/3] pipeline for provenance --- .../pipelines/data/analyses/base_analyzer.py | 479 +++++++++ .../alab/pipelines/data/config/README.md | 211 ++++ .../alab/pipelines/data/config/analyses.yaml | 135 +++ .../pipelines/data/config/config_loader.py | 327 +++++++ .../alab/pipelines/data/config/defaults.yaml | 94 ++ .../alab/pipelines/data/config/filters.yaml | 97 ++ .../alab/pipelines/data/mongodb_to_parquet.py | 591 +++++++++++ .../data/pipeline/PIPELINE_ARCHITECTURE.md | 472 +++++++++ .../alab/pipelines/data/pipeline/README.md | 267 +++++ .../alab/pipelines/data/pipeline/__init__.py | 18 + .../data/pipeline/mpcontribs_setup.py | 334 +++++++ .../pipelines/data/pipeline/pipeline_state.py | 292 ++++++ .../data/pipeline/product_pipeline.py | 691 +++++++++++++ .../pipelines/data/pipeline/run_pipeline.py | 429 ++++++++ .../pipelines/data/pipeline/s3_uploader.py | 420 ++++++++ .../data/pipeline/test_integrated_pipeline.py | 920 ++++++++++++++++++ .../pipelines/data/products/SCHEMA_SYSTEM.md | 202 ++++ .../pipelines/data/products/base_product.py | 771 +++++++++++++++ .../pipelines/data/products/schema_manager.py | 301 ++++++ .../data/products/schema_validator.py | 217 +++++ .../alab/pipelines/run_product_pipeline.sh | 38 + 21 files changed, 7306 insertions(+) create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/analyses/base_analyzer.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/README.md create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/analyses.yaml create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/config_loader.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/defaults.yaml create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/filters.yaml create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/mongodb_to_parquet.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/PIPELINE_ARCHITECTURE.md create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/README.md create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/__init__.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/mpcontribs_setup.py create mode 100755 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/pipeline_state.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/product_pipeline.py create mode 100755 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/run_pipeline.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/s3_uploader.py create mode 100755 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/test_integrated_pipeline.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/SCHEMA_SYSTEM.md create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/base_product.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/schema_manager.py create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/schema_validator.py create mode 100755 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/run_product_pipeline.sh diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/analyses/base_analyzer.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/analyses/base_analyzer.py new file mode 100644 index 000000000..72d92468b --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/analyses/base_analyzer.py @@ -0,0 +1,479 @@ +#!/usr/bin/env python3 +""" +Base Analysis Plugin Interface (Auto-Discovery) + +All analysis modules should inherit from BaseAnalyzer and implement: +1. analyze() - Run the analysis on experiments +2. get_output_schema() - Define output schema for validation + +EXTENSION: To add a new analysis: +1. Create a new file: data/analyses/your_analysis.py +2. Define a class inheriting from BaseAnalyzer +3. Set class attributes: + - name: str = "your_analysis" (used in product config) + - description: str = "What it does" + - cli_flag: str = "--your-analysis" (optional, for CLI) +4. The analysis will be auto-discovered on next pipeline run + +Analysis modules are auto-discovered from data/analyses/*.py +""" + +from abc import ABC, abstractmethod +from pathlib import Path +from typing import Dict, List, Optional, Any, Type +import pandas as pd +import logging +import importlib.util +import inspect + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +# Directory containing analysis plugins +ANALYSES_DIR = Path(__file__).parent + +# Files to skip during auto-discovery +SKIP_FILES = {'__init__.py', 'base_analyzer.py', '__pycache__'} + + +class BaseAnalyzer(ABC): + """ + Base class for all analysis modules. + + To create a new analyzer: + 1. Inherit from BaseAnalyzer + 2. Set class attributes: name, description (optional: cli_flag) + 3. Implement analyze() and get_output_schema() + + Example: + class SEMAnalyzer(BaseAnalyzer): + name = "sem_clustering" + description = "Cluster SEM images by morphology" + cli_flag = "--sem" + + def analyze(self, experiments_df, parquet_dir): + # Your analysis logic + return results_df + + def get_output_schema(self): + return {'cluster_id': {'type': 'int', 'required': True}} + """ + + # Class attributes for discovery (override in subclasses) + name: str = None # Required: unique identifier for this analyzer + description: str = "" # Optional: human-readable description + cli_flag: str = None # Optional: command line flag (e.g., "--xrd") + + def __init__(self, config: Dict = None): + """ + Initialize analyzer with configuration + + Args: + config: Analysis-specific configuration from product config + """ + self.config = config or {} + # Use class name attribute or fall back to class name + if self.name is None: + self.name = self.__class__.__name__.lower().replace('analyzer', '') + + @abstractmethod + def analyze(self, + experiments_df: pd.DataFrame, + parquet_dir: Path) -> pd.DataFrame: + """ + Run analysis on experiments + + Args: + experiments_df: DataFrame of experiments to analyze + parquet_dir: Directory containing parquet files + + Returns: + DataFrame with analysis results (one row per experiment) + """ + pass + + @abstractmethod + def get_output_schema(self) -> Dict[str, Dict]: + """ + Get output schema for validation + + Returns: + Dict of field_name -> {type, description, required} + """ + pass + + def validate_output(self, results_df: pd.DataFrame) -> bool: + """ + Validate that results match expected schema + + Args: + results_df: Analysis results DataFrame + + Returns: + True if valid + """ + schema = self.get_output_schema() + + for field, spec in schema.items(): + if spec.get('required', False) and field not in results_df.columns: + logger.error(f"Missing required field: {field}") + return False + + if field in results_df.columns: + # Type checking could be added here + pass + + return True + + def save_results(self, results_df: pd.DataFrame, output_dir: Path): + """Save analysis results to parquet""" + output_dir.mkdir(parents=True, exist_ok=True) + output_file = output_dir / f"{self.name.lower()}_results.parquet" + results_df.to_parquet(output_file, index=False) + logger.info(f"Saved {len(results_df)} results to {output_file}") + + +class XRDAnalyzer(BaseAnalyzer): + """XRD Phase Analysis using DARA""" + + # Class attributes for discovery + name = "xrd_dara" + description = "XRD phase identification using DARA (Deep Analysis for Rietveld Automation)" + cli_flag = "--xrd" + + def __init__(self, config: Dict = None): + super().__init__(config) + + def analyze(self, experiments_df: pd.DataFrame, parquet_dir: Path) -> pd.DataFrame: + """Run DARA XRD analysis on experiments""" + import subprocess + import json + + results = [] + xrd_results_dir = Path("data/xrd_creation/results") + + for _, exp in experiments_df.iterrows(): + exp_name = exp['name'] + result_file = xrd_results_dir / f"{exp_name}_result.json" + + # Check if already analyzed + if result_file.exists(): + with open(result_file, 'r') as f: + result = json.load(f) + results.append({ + 'experiment_name': exp_name, + 'xrd_success': result.get('success', False), + 'xrd_rwp': result.get('rwp'), + 'xrd_num_phases': result.get('num_phases', 0), + 'xrd_error': result.get('error') + }) + else: + # Run analysis + try: + cmd = [ + "python", + "data/xrd_creation/analyze_single.py", + "--experiment", exp_name, + "--mode", "phase_search" + ] + + # Add config options + if self.config.get('wmin'): + cmd.extend(["--wmin", str(self.config['wmin'])]) + if self.config.get('wmax'): + cmd.extend(["--wmax", str(self.config['wmax'])]) + + result = subprocess.run(cmd, capture_output=True, text=True, timeout=120) + + # Load result + if result_file.exists(): + with open(result_file, 'r') as f: + analysis_result = json.load(f) + results.append({ + 'experiment_name': exp_name, + 'xrd_success': analysis_result.get('success', False), + 'xrd_rwp': analysis_result.get('rwp'), + 'xrd_num_phases': analysis_result.get('num_phases', 0), + 'xrd_error': analysis_result.get('error') + }) + else: + results.append({ + 'experiment_name': exp_name, + 'xrd_success': False, + 'xrd_rwp': None, + 'xrd_num_phases': 0, + 'xrd_error': 'Analysis failed' + }) + + except Exception as e: + logger.error(f"XRD analysis failed for {exp_name}: {e}") + results.append({ + 'experiment_name': exp_name, + 'xrd_success': False, + 'xrd_rwp': None, + 'xrd_num_phases': 0, + 'xrd_error': str(e) + }) + + return pd.DataFrame(results) + + def get_output_schema(self) -> Dict[str, Dict]: + return { + 'xrd_success': {'type': 'boolean', 'required': True, 'description': 'XRD analysis succeeded'}, + 'xrd_rwp': {'type': 'float', 'required': False, 'description': 'Weighted profile R-factor'}, + 'xrd_num_phases': {'type': 'int', 'required': True, 'description': 'Number of phases identified'}, + 'xrd_error': {'type': 'string', 'required': False, 'description': 'Error message if failed'} + } + + +class PowderStatisticsAnalyzer(BaseAnalyzer): + """Analyze powder dosing statistics""" + + # Class attributes for discovery + name = "powder_statistics" + description = "Calculate powder dosing accuracy and statistics" + cli_flag = "--powder-stats" + + def __init__(self, config: Dict = None): + super().__init__(config) + + def analyze(self, experiments_df: pd.DataFrame, parquet_dir: Path) -> pd.DataFrame: + """Calculate powder dosing statistics""" + + # Load powder dosing data + powder_doses_df = pd.read_parquet(parquet_dir / "powder_doses.parquet") + + results = [] + + for _, exp in experiments_df.iterrows(): + exp_id = exp['experiment_id'] + exp_name = exp['name'] + + # Get doses for this experiment + exp_doses = powder_doses_df[powder_doses_df['experiment_id'] == exp_id] + + if len(exp_doses) > 0: + # Calculate statistics + avg_accuracy = exp_doses['accuracy_percent'].mean() if 'accuracy_percent' in exp_doses.columns else None + total_doses = len(exp_doses) + unique_powders = exp_doses['powder_name'].nunique() + total_mass = exp_doses['actual_mass'].sum() + + results.append({ + 'experiment_name': exp_name, + 'powder_avg_accuracy': avg_accuracy, + 'powder_total_doses': total_doses, + 'powder_unique_count': unique_powders, + 'powder_total_mass_g': total_mass + }) + else: + results.append({ + 'experiment_name': exp_name, + 'powder_avg_accuracy': None, + 'powder_total_doses': 0, + 'powder_unique_count': 0, + 'powder_total_mass_g': 0.0 + }) + + return pd.DataFrame(results) + + def get_output_schema(self) -> Dict[str, Dict]: + return { + 'powder_avg_accuracy': {'type': 'float', 'required': False, 'description': 'Average dosing accuracy %'}, + 'powder_total_doses': {'type': 'int', 'required': True, 'description': 'Total number of doses'}, + 'powder_unique_count': {'type': 'int', 'required': True, 'description': 'Number of unique powders'}, + 'powder_total_mass_g': {'type': 'float', 'required': True, 'description': 'Total powder mass in grams'} + } + + +class AnalysisPluginManager: + """ + Manages and auto-discovers analysis plugins. + + Analyzers are discovered from: + 1. Built-in analyzers in this file (XRDAnalyzer, PowderStatisticsAnalyzer) + 2. Any .py file in data/analyses/ containing a BaseAnalyzer subclass + + To add a new analyzer: + 1. Create data/analyses/your_analyzer.py + 2. Define class YourAnalyzer(BaseAnalyzer) with name attribute + 3. It will be auto-discovered + """ + + def __init__(self, analyses_dir: Path = None): + self.analyses_dir = Path(analyses_dir or ANALYSES_DIR) + self.analyzers: Dict[str, Type[BaseAnalyzer]] = {} + self._discover_analyzers() + + def _discover_analyzers(self): + """ + Auto-discover available analyzers from the analyses directory. + + Discovery rules: + 1. Register built-in analyzers first + 2. Scan all .py files in analyses_dir (except SKIP_FILES) + 3. Find classes that inherit from BaseAnalyzer + 4. Register using the class's 'name' attribute + """ + # Built-in analyzers (defined in this file) + self.analyzers['xrd_dara'] = XRDAnalyzer + self.analyzers['powder_statistics'] = PowderStatisticsAnalyzer + + # Auto-discover from analyses directory + if not self.analyses_dir.exists(): + logger.warning(f"Analyses directory not found: {self.analyses_dir}") + return + + for analyzer_file in sorted(self.analyses_dir.glob("*.py")): + if analyzer_file.name in SKIP_FILES: + continue + + try: + # Import the module + module_name = analyzer_file.stem + spec = importlib.util.spec_from_file_location(module_name, analyzer_file) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + # Find BaseAnalyzer subclasses + for name, obj in inspect.getmembers(module, inspect.isclass): + # Skip imported classes and BaseAnalyzer itself + if obj.__module__ != module.__name__: + continue + if not issubclass(obj, BaseAnalyzer) or obj is BaseAnalyzer: + continue + + # Get analyzer name from class attribute + analyzer_name = getattr(obj, 'name', None) + if analyzer_name is None: + # Derive from class name: SEMAnalyzer -> sem + analyzer_name = name.lower().replace('analyzer', '') + + # Don't override built-in analyzers + if analyzer_name not in self.analyzers: + self.analyzers[analyzer_name] = obj + logger.debug(f"Discovered analyzer: {analyzer_name} -> {name}") + + except Exception as e: + logger.error(f"Failed to load analyzer from {analyzer_file}: {e}") + + logger.info(f"Discovered {len(self.analyzers)} analyzers: {', '.join(self.analyzers.keys())}") + + def get_analyzer(self, name: str, config: Dict = None) -> Optional[BaseAnalyzer]: + """Get analyzer instance by name""" + if name in self.analyzers: + return self.analyzers[name](config) + + logger.warning(f"Analyzer '{name}' not found. Available: {', '.join(self.analyzers.keys())}") + return None + + def list_analyzers(self) -> List[str]: + """List available analyzer names""" + return sorted(self.analyzers.keys()) + + def get_analyzer_info(self) -> Dict[str, Dict]: + """Get info about all available analyzers""" + info = {} + for name, analyzer_class in self.analyzers.items(): + info[name] = { + 'class': analyzer_class.__name__, + 'description': getattr(analyzer_class, 'description', ''), + 'cli_flag': getattr(analyzer_class, 'cli_flag', None), + } + return info + + def run_analyses(self, + analyses: List[Dict], + experiments_df: pd.DataFrame, + parquet_dir: Path, + output_dir: Path) -> pd.DataFrame: + """ + Run multiple analyses and merge results + + Args: + analyses: List of {name, config} dicts + experiments_df: Experiments to analyze + parquet_dir: Input parquet directory + output_dir: Output directory for results + + Returns: + Merged DataFrame with all analysis results + """ + # Start with essential experiment fields + base_fields = ['experiment_id', 'name', 'experiment_type', 'target_formula'] + # Only include fields that exist in the DataFrame + available_fields = [f for f in base_fields if f in experiments_df.columns] + all_results = experiments_df[available_fields].copy() + + for analysis_config in analyses: + if not analysis_config.get('enabled', True): + continue + + analyzer = self.get_analyzer( + analysis_config['name'], + analysis_config.get('config', {}) + ) + + if analyzer: + logger.info(f"Running {analyzer.name} analysis...") + + try: + results = analyzer.analyze(experiments_df, parquet_dir) + + if analyzer.validate_output(results): + # Merge results + all_results = all_results.merge( + results, + left_on='name', + right_on='experiment_name', + how='left' + ) + + # Save individual results + analyzer.save_results(results, output_dir) + else: + logger.error(f"Validation failed for {analyzer.name}") + + except Exception as e: + logger.error(f"Analysis {analyzer.name} failed: {e}") + + return all_results + + +# CLI for listing available analyzers +if __name__ == '__main__': + from rich.console import Console + from rich.table import Table + + console = Console() + + manager = AnalysisPluginManager() + + console.print(f"\n[bold cyan]Analysis Plugin Auto-Discovery[/bold cyan]") + console.print(f"Directory: {manager.analyses_dir}") + console.print(f"Discovered: {len(manager.analyzers)} analyzers\n") + + # Create table + table = Table(title="Available Analyzers") + table.add_column("Name", style="cyan") + table.add_column("Class", style="green") + table.add_column("CLI Flag", style="yellow") + table.add_column("Description") + + for name, info in manager.get_analyzer_info().items(): + table.add_row( + name, + info['class'], + info['cli_flag'] or "-", + info['description'] or "-" + ) + + console.print(table) + + console.print("\n[bold green]✓ Auto-discovery complete![/bold green]") + console.print("\n[dim]To add a new analyzer:[/dim]") + console.print(" 1. Create: data/analyses/your_analyzer.py") + console.print(" 2. Define: class YourAnalyzer(BaseAnalyzer): ...") + console.print(" 3. Set: name = 'your_analyzer'") + console.print(" 4. Implement: analyze() and get_output_schema()") + console.print(" 5. Run pipeline - analyzer auto-discovered!") diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/README.md b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/README.md new file mode 100644 index 000000000..e7a8b7005 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/README.md @@ -0,0 +1,211 @@ +# A-Lab Pipeline Configuration + +Configuration system with three-layer priority: + +``` +1. Environment Variables (highest) → 2. YAML Files → 3. Code Defaults (fallback) +``` + +## Quick Start + +### Using Environment Variables (Recommended for Production) + +**Option 1: Copy example file** + +```bash +# Copy the example .env file (with current defaults) +cp data/config/env.example .env + +# Edit .env and uncomment/modify values +# Then source it before running +source .env +./update_data.sh +``` + +**Option 2: Export directly** + +```bash +# Set MongoDB connection +export ALAB_MONGO_URI="mongodb://production-host:27017/" +export ALAB_MONGO_DB="production" + +# Set S3 bucket +export ALAB_S3_BUCKET="my-custom-bucket" + +# Run pipeline (uses env vars automatically) +./update_data.sh +``` + +### Using YAML Files (Recommended for Development) + +Edit `data/config/defaults.yaml`: + +```yaml +mongodb: + uri: 'mongodb://localhost:27017/' + database: 'my_database' + collection: 'my_collection' +``` + +### View Current Configuration + +```bash +python data/config/config_loader.py +``` + +Shows all loaded values and their sources (env vs yaml vs defaults). + +## Configuration Files + +| File | Purpose | +| -------------------- | ------------------------------- | +| **defaults.yaml** | Global pipeline defaults | +| **filters.yaml** | Experiment filter presets | +| **analyses.yaml** | Analysis plugin documentation | +| **config_loader.py** | Configuration loading system | +| **env.example** | Example env file (copy to .env) | + +## Environment Variables + +### MongoDB + +```bash +ALAB_MONGO_URI=mongodb://localhost:27017/ # MongoDB connection URI +ALAB_MONGO_DB=temporary # Database name +ALAB_MONGO_COLLECTION=release # Collection name +``` + +### S3 Upload + +```bash +ALAB_S3_BUCKET=materialsproject-contribs # S3 bucket name +ALAB_S3_PREFIX=alab_synthesis # S3 prefix path +ALAB_S3_EXCLUDE_LARGE=true # Exclude large files +ALAB_S3_LARGE_THRESHOLD_MB=50 # Large file threshold (MB) +``` + +### Parquet Options + +```bash +ALAB_SKIP_TEMP_LOGS=false # Skip temperature logs +ALAB_SKIP_XRD_POINTS=false # Skip XRD data points +ALAB_SKIP_WORKFLOW_TASKS=false # Skip workflow tasks +ALAB_PARQUET_COMPRESSION=snappy # Compression: snappy, gzip, brotli +ALAB_PARQUET_ENGINE=pyarrow # Engine: pyarrow, fastparquet +``` + +### Materials Project API + +```bash +ALAB_MP_API_KEY=your_api_key # MP API key (for XRD analysis) +# OR +MP_API_KEY=your_api_key # Alternative name +``` + +## Usage in Scripts + +### Python + +```python +from config_loader import get_config + +# Get configuration +config = get_config() + +# Access values +print(config.mongo_uri) # mongodb://localhost:27017/ +print(config.mongo_db) # temporary +print(config.s3_bucket) # materialsproject-contribs + +# Or use convenience functions +from config_loader import get_mongo_uri, get_s3_bucket + +uri = get_mongo_uri() # Gets from env > yaml > default +bucket = get_s3_bucket() +``` + +### Shell Scripts + +```bash +# Use environment variables directly +: ${ALAB_MONGO_URI:="mongodb://localhost:27017/"} + +# Or source from .env file +if [ -f data/.env ]; then + export $(grep -v '^#' data/.env | xargs) +fi +``` + +## Configuration Priority Examples + +### Example 1: All from YAML + +```bash +# No env vars set +$ python data/config/config_loader.py +MongoDB URI: mongodb://localhost:27017/ (from YAML) +``` + +### Example 2: Override with Env + +```bash +# Set env var +$ export ALAB_MONGO_URI="mongodb://production:27017/" +$ python data/config/config_loader.py +MongoDB URI: mongodb://production:27017/ (from ENV) ✓ +``` + +### Example 3: Mixed Sources + +```bash +# Some from env, some from yaml +$ export ALAB_MONGO_URI="mongodb://prod:27017/" # Custom URI +# Leave ALAB_MONGO_DB unset # Use YAML default +$ python data/config/config_loader.py +MongoDB URI: mongodb://prod:27017/ (from ENV) ✓ +MongoDB DB: temporary (from YAML) +``` + +## Best Practices + +1. **Development**: Use `defaults.yaml` for local development +2. **Production**: Use environment variables for sensitive values +3. **Testing**: Use env vars to point to test databases +4. **CI/CD**: Set env vars in your deployment pipeline +5. **Never commit** `.env` files (already in `.gitignore`) + +## Troubleshooting + +### Config not loading? + +```bash +# Check current config +python data/config/config_loader.py + +# Verify env vars are set +env | grep ALAB_ +``` + +### Want to use .env file? + +```bash +# Create from example +cp data/config/env.example .env + +# Edit .env with your values (uncomment lines to override defaults) +nano .env + +# Source it before running scripts +source .env +./update_data.sh +``` + +### Reset to defaults + +```bash +# Unset all ALAB env vars +unset $(env | grep ALAB_ | cut -d= -f1) + +# Now uses YAML/defaults only +./update_data.sh +``` diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/analyses.yaml b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/analyses.yaml new file mode 100644 index 000000000..5290a7a3f --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/analyses.yaml @@ -0,0 +1,135 @@ +# ============================================================================= +# A-Lab Analysis Registry +# ============================================================================= +# Documentation for available analysis plugins. +# Analyses are auto-discovered from data/analyses/*.py +# +# This file serves as: +# 1. Documentation of available analyses +# 2. Default configuration for each analysis +# 3. Template for adding new analyses +# ============================================================================= + +# Built-in analyses (always available) +analyses: + xrd_dara: + description: 'XRD phase identification using DARA' + class: XRDAnalyzer + file: base_analyzer.py + cli_flag: '--xrd' + output_parquet: xrd_refinements.parquet, xrd_phases.parquet + default_config: + wmin: 10 # Minimum 2-theta angle + wmax: 80 # Maximum 2-theta angle + save_viz: false # Save visualization images + outputs: + - xrd_success: 'Whether analysis succeeded' + - xrd_rwp: 'Weighted profile R-factor' + - xrd_num_phases: 'Number of phases identified' + - xrd_error: 'Error message if failed' + requirements: + - experiments.parquet (with xrd_sampleid_in_aeris) + - xrd_data_points.parquet (optional, for patterns) + notes: | + Uses DARA (Deep Analysis for Rietveld Automation) for automated + phase identification. Requires MP API key for CIF downloads. + Results stored in data/xrd_creation/results/ + + powder_statistics: + description: 'Calculate powder dosing statistics' + class: PowderStatisticsAnalyzer + file: base_analyzer.py + cli_flag: '--powder-stats' + output_parquet: null # Results merged into main output + default_config: {} + outputs: + - powder_avg_accuracy: 'Average dosing accuracy %' + - powder_total_doses: 'Total number of doses' + - powder_unique_count: 'Number of unique powders' + - powder_total_mass_g: 'Total powder mass in grams' + requirements: + - experiments.parquet + - powder_doses.parquet + notes: | + Calculates statistics about powder dosing accuracy. + Fast analysis, recommended for all products. + +# ============================================================================= +# Adding a New Analysis +# ============================================================================= +# To add a new analysis (e.g., SEM clustering): +# +# 1. Create the analyzer file: +# data/analyses/sem_analyzer.py +# +# 2. Define the analyzer class: +# ```python +# from base_analyzer import BaseAnalyzer +# +# class SEMAnalyzer(BaseAnalyzer): +# name = "sem_clustering" +# description = "Cluster SEM images by morphology" +# cli_flag = "--sem" +# +# def analyze(self, experiments_df, parquet_dir): +# # Your analysis logic here +# results = [] +# for _, exp in experiments_df.iterrows(): +# # Process each experiment +# results.append({ +# 'experiment_name': exp['name'], +# 'cluster_id': compute_cluster(exp), +# 'morphology_score': compute_score(exp) +# }) +# return pd.DataFrame(results) +# +# def get_output_schema(self): +# return { +# 'cluster_id': {'type': 'int', 'required': True}, +# 'morphology_score': {'type': 'float', 'required': False} +# } +# ``` +# +# 3. Document here (optional): +# sem_clustering: +# description: "Cluster SEM images by morphology" +# class: SEMAnalyzer +# file: sem_analyzer.py +# ... +# +# 4. The analysis will be auto-discovered on next pipeline run +# ============================================================================= + +# Placeholder for future analyses +# Uncomment and modify when adding: + +# sem_clustering: +# description: "Cluster SEM images by morphology" +# class: SEMAnalyzer +# file: sem_analyzer.py +# cli_flag: "--sem" +# default_config: +# num_clusters: 5 +# feature_extraction: "resnet50" +# outputs: +# - cluster_id: "Cluster assignment" +# - morphology_score: "Morphology similarity score" +# requirements: +# - SEM images in experiments/*/SEM images/ +# notes: | +# Uses computer vision to cluster SEM images. +# Requires tensorflow/pytorch. + +# heating_profile: +# description: "Analyze heating profile characteristics" +# class: HeatingProfileAnalyzer +# file: heating_analyzer.py +# cli_flag: "--heating" +# default_config: {} +# outputs: +# - heating_rate_avg: "Average heating rate" +# - overshoot_celsius: "Temperature overshoot" +# - time_at_target: "Time at target temperature" +# requirements: +# - temperature_logs.parquet + diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/config_loader.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/config_loader.py new file mode 100644 index 000000000..5d6163017 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/config_loader.py @@ -0,0 +1,327 @@ +#!/usr/bin/env python3 +""" +Configuration Loader for A-Lab Pipeline + +Loads configuration with priority: +1. Environment variables (highest priority) +2. YAML config files +3. Code defaults (fallback) + +Environment Variables: + ALAB_MONGO_URI - MongoDB connection URI + ALAB_MONGO_DB - MongoDB database name + ALAB_MONGO_COLLECTION - MongoDB collection name + ALAB_S3_BUCKET - S3 bucket for uploads + ALAB_S3_PREFIX - S3 prefix path + ALAB_MP_API_KEY - Materials Project API key (for XRD analysis) + +Example: + export ALAB_MONGO_URI="mongodb://production:27017/" + export ALAB_S3_BUCKET="my-custom-bucket" +""" + +import os +import yaml +from pathlib import Path +from typing import Dict, Any, Optional +from functools import lru_cache + +# Config file paths +CONFIG_DIR = Path(__file__).parent +DEFAULTS_PATH = CONFIG_DIR / "defaults.yaml" + + +class ConfigLoader: + """Load configuration from env vars -> YAML -> defaults""" + + def __init__(self): + self._defaults = None + self._load_defaults() + + def _load_defaults(self): + """Load defaults.yaml file""" + if DEFAULTS_PATH.exists(): + with open(DEFAULTS_PATH, 'r') as f: + self._defaults = yaml.safe_load(f) or {} + else: + self._defaults = {} + + def _get_env_or_yaml(self, env_var: str, yaml_path: list, default: Any = None) -> Any: + """ + Get value from env var first, then YAML, then default. + + Args: + env_var: Environment variable name + yaml_path: Path through YAML dict (e.g., ['mongodb', 'uri']) + default: Fallback default value + + Returns: + Configuration value + """ + # 1. Try environment variable + env_value = os.getenv(env_var) + if env_value is not None: + return env_value + + # 2. Try YAML config + value = self._defaults + for key in yaml_path: + if isinstance(value, dict) and key in value: + value = value[key] + else: + value = None + break + + if value is not None: + return value + + # 3. Use default + return default + + # MongoDB configuration + @property + def mongo_uri(self) -> str: + """MongoDB connection URI""" + return self._get_env_or_yaml( + 'ALAB_MONGO_URI', + ['mongodb', 'uri'], + 'mongodb://localhost:27017/' + ) + + @property + def mongo_db(self) -> str: + """MongoDB database name""" + return self._get_env_or_yaml( + 'ALAB_MONGO_DB', + ['mongodb', 'database'], + 'temporary' + ) + + @property + def mongo_collection(self) -> str: + """MongoDB collection name""" + return self._get_env_or_yaml( + 'ALAB_MONGO_COLLECTION', + ['mongodb', 'collection'], + 'release' + ) + + # Parquet configuration + @property + def parquet_skip_temp_logs(self) -> bool: + """Skip temperature logs during transformation""" + value = self._get_env_or_yaml( + 'ALAB_SKIP_TEMP_LOGS', + ['parquet', 'skip_temperature_logs'], + False + ) + return self._to_bool(value) + + @property + def parquet_skip_xrd_points(self) -> bool: + """Skip XRD data points during transformation""" + value = self._get_env_or_yaml( + 'ALAB_SKIP_XRD_POINTS', + ['parquet', 'skip_xrd_points'], + False + ) + return self._to_bool(value) + + @property + def parquet_skip_workflow_tasks(self) -> bool: + """Skip workflow tasks during transformation""" + value = self._get_env_or_yaml( + 'ALAB_SKIP_WORKFLOW_TASKS', + ['parquet', 'skip_workflow_tasks'], + False + ) + return self._to_bool(value) + + @property + def parquet_compression(self) -> str: + """Parquet compression type""" + return self._get_env_or_yaml( + 'ALAB_PARQUET_COMPRESSION', + ['parquet', 'compression'], + 'snappy' + ) + + @property + def parquet_engine(self) -> str: + """Parquet engine""" + return self._get_env_or_yaml( + 'ALAB_PARQUET_ENGINE', + ['parquet', 'engine'], + 'pyarrow' + ) + + # S3 Upload configuration + @property + def s3_bucket(self) -> str: + """S3 bucket for uploads""" + return self._get_env_or_yaml( + 'ALAB_S3_BUCKET', + ['upload', 's3_bucket'], + 'materialsproject-contribs' + ) + + @property + def s3_prefix(self) -> str: + """S3 prefix path""" + return self._get_env_or_yaml( + 'ALAB_S3_PREFIX', + ['upload', 's3_prefix'], + 'alab_synthesis' + ) + + @property + def s3_exclude_large_files(self) -> bool: + """Exclude large files from S3 upload""" + value = self._get_env_or_yaml( + 'ALAB_S3_EXCLUDE_LARGE', + ['upload', 'exclude_large_files'], + True + ) + return self._to_bool(value) + + @property + def s3_large_file_threshold_mb(self) -> int: + """Large file threshold in MB""" + value = self._get_env_or_yaml( + 'ALAB_S3_LARGE_THRESHOLD_MB', + ['upload', 'large_file_threshold_mb'], + 50 + ) + return int(value) + + # Materials Project API + @property + def mp_api_key(self) -> Optional[str]: + """Materials Project API key (for XRD analysis)""" + return os.getenv('ALAB_MP_API_KEY') or os.getenv('MP_API_KEY') + + # Utility methods + @staticmethod + def _to_bool(value: Any) -> bool: + """Convert various types to boolean""" + if isinstance(value, bool): + return value + if isinstance(value, str): + return value.lower() in ('true', '1', 'yes', 'on') + return bool(value) + + def get_all(self) -> Dict[str, Any]: + """Get all configuration as dict""" + return { + 'mongodb': { + 'uri': self.mongo_uri, + 'database': self.mongo_db, + 'collection': self.mongo_collection, + }, + 'parquet': { + 'skip_temperature_logs': self.parquet_skip_temp_logs, + 'skip_xrd_points': self.parquet_skip_xrd_points, + 'skip_workflow_tasks': self.parquet_skip_workflow_tasks, + 'compression': self.parquet_compression, + 'engine': self.parquet_engine, + }, + 's3': { + 'bucket': self.s3_bucket, + 'prefix': self.s3_prefix, + 'exclude_large_files': self.s3_exclude_large_files, + 'large_file_threshold_mb': self.s3_large_file_threshold_mb, + }, + 'mp_api_key': self.mp_api_key, + } + + +# Singleton instance +@lru_cache(maxsize=1) +def get_config() -> ConfigLoader: + """Get cached configuration loader instance""" + return ConfigLoader() + + +# Convenience functions for direct access +def get_mongo_uri() -> str: + """Get MongoDB URI from env > yaml > default""" + return get_config().mongo_uri + + +def get_mongo_db() -> str: + """Get MongoDB database from env > yaml > default""" + return get_config().mongo_db + + +def get_mongo_collection() -> str: + """Get MongoDB collection from env > yaml > default""" + return get_config().mongo_collection + + +def get_s3_bucket() -> str: + """Get S3 bucket from env > yaml > default""" + return get_config().s3_bucket + + +def get_s3_prefix() -> str: + """Get S3 prefix from env > yaml > default""" + return get_config().s3_prefix + + +if __name__ == '__main__': + # Display current configuration + config = get_config() + + print("=" * 60) + print("A-Lab Pipeline Configuration") + print("=" * 60) + print("\nConfiguration Priority: ENV → YAML → Defaults\n") + + print("MongoDB:") + print(f" URI: {config.mongo_uri}") + print(f" Database: {config.mongo_db}") + print(f" Collection: {config.mongo_collection}") + + print("\nParquet:") + print(f" Skip temp logs: {config.parquet_skip_temp_logs}") + print(f" Skip XRD points: {config.parquet_skip_xrd_points}") + print(f" Skip workflow tasks: {config.parquet_skip_workflow_tasks}") + print(f" Compression: {config.parquet_compression}") + print(f" Engine: {config.parquet_engine}") + + print("\nS3 Upload:") + print(f" Bucket: {config.s3_bucket}") + print(f" Prefix: {config.s3_prefix}") + print(f" Exclude large files: {config.s3_exclude_large_files}") + print(f" Large file threshold: {config.s3_large_file_threshold_mb} MB") + + print("\nMaterials Project:") + mp_key = config.mp_api_key + if mp_key: + print(f" API Key: {mp_key[:10]}... (found)") + else: + print(f" API Key: Not set") + + print("\n" + "=" * 60) + + # Show which values are from env + print("\nEnvironment Variables Detected:") + env_vars = [ + 'ALAB_MONGO_URI', 'ALAB_MONGO_DB', 'ALAB_MONGO_COLLECTION', + 'ALAB_S3_BUCKET', 'ALAB_S3_PREFIX', + 'ALAB_SKIP_TEMP_LOGS', 'ALAB_SKIP_XRD_POINTS', + 'ALAB_MP_API_KEY', 'MP_API_KEY' + ] + + found_envs = [] + for var in env_vars: + if os.getenv(var): + found_envs.append(f" ✓ {var}") + + if found_envs: + print("\n".join(found_envs)) + else: + print(" (none) - using YAML/defaults") + + print("\n" + "=" * 60) + diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/defaults.yaml b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/defaults.yaml new file mode 100644 index 000000000..da34782ad --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/defaults.yaml @@ -0,0 +1,94 @@ +# ============================================================================= +# A-Lab Pipeline Default Configuration +# ============================================================================= +# This file contains global defaults for the pipeline. +# Product-specific configs can override these values. +# +# Configuration Priority: +# 1. Environment variables (highest - for deployment/secrets) +# 2. This YAML file (middle - for general config) +# 3. Code defaults (lowest - fallback only) +# +# Environment Variables (see .env.example for full list): +# ALAB_MONGO_URI, ALAB_MONGO_DB, ALAB_MONGO_COLLECTION +# ALAB_S3_BUCKET, ALAB_S3_PREFIX +# ALAB_MP_API_KEY +# +# To customize: Either set env vars OR copy values to your product's config.yaml +# ============================================================================= + +# Pipeline version +version: '2.0' + +# MongoDB connection (used by mongodb_to_parquet.py) +# Override with: ALAB_MONGO_URI, ALAB_MONGO_DB, ALAB_MONGO_COLLECTION +mongodb: + uri: 'mongodb://localhost:27017/' + database: 'temporary' + collection: 'release' + +# Parquet transformation defaults +# Override with: ALAB_SKIP_TEMP_LOGS, ALAB_SKIP_XRD_POINTS, etc. +parquet: + # Skip large arrays by default for faster processing + skip_temperature_logs: false + skip_xrd_points: false + skip_workflow_tasks: false + # Compression settings + compression: 'snappy' + engine: 'pyarrow' + +# Default experiment filter (can be overridden per product) +experiment_filter: + # Status filter (completed experiments by default) + status: + - completed + # Require XRD data by default + has_xrd: true + # No date range by default (all dates) + date_range: null + +# Analysis defaults +analyses: + # XRD analysis settings + xrd_dara: + enabled: true + config: + wmin: 10 + wmax: 80 + save_viz: false + + # Powder statistics (always fast, recommend enabling) + powder_statistics: + enabled: true + config: {} + +# Upload settings +# Override with: ALAB_S3_BUCKET, ALAB_S3_PREFIX, ALAB_S3_EXCLUDE_LARGE +upload: + # S3 bucket for OpenData uploads + s3_bucket: 'materialsproject-contribs' + s3_prefix: 'alab_synthesis' + # Skip large files by default + exclude_large_files: true + large_file_threshold_mb: 50 + +# Diagram generation +diagram: + # Generate diagram after parquet creation + auto_generate: true + # Output formats + formats: + - summary # Markdown summary + - mermaid # ERD diagram + # Output filename (relative to product directory) + output_file: 'SCHEMA_DIAGRAM.md' + +# Validation settings +validation: + # Sample size for row-by-row validation (full validation can be slow) + sample_size: 100 + # Fail pipeline on validation errors + fail_on_error: false + # Show first N errors + max_errors_shown: 10 diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/filters.yaml b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/filters.yaml new file mode 100644 index 000000000..34de32236 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/filters.yaml @@ -0,0 +1,97 @@ +# ============================================================================= +# A-Lab Experiment Filter Presets +# ============================================================================= +# Predefined filters for common experiment selections. +# Use these in your product config or as templates. +# +# Usage in product config.yaml: +# experiment_filter: +# preset: "completed_with_xrd" # Use a preset +# OR +# experiment_filter: +# types: [NSC, Na] # Custom filter (types auto-discovered from MongoDB) +# ============================================================================= + +# NOTE: Experiment types are AUTO-DISCOVERED from MongoDB +# The pipeline discovers available types by querying MongoDB at runtime. +# Run this to see available types: +# python data/tools/analyze_mongodb.py +# +# Example types (discovered from A-Lab): +# NSC, Na, PG, MINES, TRI, etc. +# These change based on what's actually in your MongoDB database. + +# Valid experiment statuses (standardized) +statuses: + - completed # All tasks finished successfully + - error # At least one task failed + - active # Currently running + - unknown # Status could not be determined + +# ============================================================================= +# Filter Presets (Generic) +# ============================================================================= +# Generic filter presets that work regardless of experiment types in MongoDB. +# Type-specific filters are auto-generated during product creation. + +presets: + # All completed experiments with XRD data + completed_with_xrd: + description: 'Completed experiments that have XRD measurements' + status: + - completed + has_xrd: true + + # All completed experiments (with or without XRD) + all_completed: + description: 'All completed experiments' + status: + - completed + has_xrd: null # Don't filter by XRD + + # Failed experiments (for debugging) + failed: + description: 'Experiments with errors (for debugging)' + status: + - error + has_xrd: null + + # Everything (no filtering) + all: + description: 'All experiments (no filtering)' + # No filters applied +# ============================================================================= +# Type-Specific Filters (Auto-Generated) +# ============================================================================= +# Filters for specific experiment types (NSC, Na, PG, etc.) are auto-generated +# during product creation based on what's discovered in MongoDB. +# +# The pipeline will: +# 1. Query MongoDB for available experiment types +# 2. Show hierarchical selection (root types and subgroups) +# 3. Generate appropriate filters based on your selection +# +# You don't need to define type-specific presets here. +# ============================================================================= +# Custom Filters +# ============================================================================= +# Add your own generic filter presets here. +# For type-specific filters, use the interactive product creation which will +# auto-discover available types from MongoDB. +# +# Generic filter template: +# +# my_custom_filter: +# description: "What this filter selects" +# status: [completed] # Optional: status filter +# has_xrd: true # Optional: XRD requirement +# has_sem: true # Optional: SEM requirement (future) +# date_range: # Optional: date range +# start: "2024-01-01" +# end: "2024-12-31" +# +# NOTE: To filter by specific experiment types (e.g., NSC, Na): +# - Use interactive product creation (./run_product_pipeline.sh create) +# - Types are discovered from MongoDB at runtime +# - Or specify directly in product config.yaml: types: [NSC, Na] + diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/mongodb_to_parquet.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/mongodb_to_parquet.py new file mode 100644 index 000000000..92bdd7611 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/mongodb_to_parquet.py @@ -0,0 +1,591 @@ +#!/usr/bin/env python3 +""" +MongoDB to Parquet Data Pipeline +Transforms A-Lab experiment data from MongoDB to Parquet files for dashboard + +Output Structure (Consolidated): +- experiments.parquet: Main table with ALL 1:1 data merged (~45 columns) +- experiment_elements.parquet: Elements per experiment (1:N) +- powder_doses.parquet: Individual powder doses (1:N) +- temperature_logs.parquet: Temperature readings (1:N, optional) +- xrd_data_points.parquet: Raw XRD patterns (1:N, optional) +- workflow_tasks.parquet: Task execution history (1:N, optional) +- xrd_refinements.parquet: Analysis results (from DARA) +- xrd_phases.parquet: Identified phases (from DARA) +""" + +import sys +import os +import logging +from pathlib import Path +from datetime import datetime +from typing import Dict, List, Optional, Any +import argparse + +import pandas as pd +import numpy as np +from pymongo import MongoClient +from tqdm import tqdm + +# Import config loader +sys.path.insert(0, str(Path(__file__).parent / "config")) +from config_loader import get_config + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(levelname)s - %(message)s', + handlers=[ + logging.FileHandler('parquet_migration.log'), + logging.StreamHandler() + ] +) +logger = logging.getLogger(__name__) + +# Load configuration (env vars > yaml > defaults) +config = get_config() + +# Output directory for parquet files +OUTPUT_DIR = Path(__file__).parent / "parquet" + + +class MongoToParquetTransformer: + """Transform MongoDB experiments to consolidated Parquet files""" + + def __init__(self, mongo_uri: str = None, mongo_db: str = None, + mongo_collection: str = None, output_dir: Path = OUTPUT_DIR): + # Use config if not explicitly provided + self.mongo_uri = mongo_uri or config.mongo_uri + self.mongo_db = mongo_db or config.mongo_db + self.mongo_collection = mongo_collection or config.mongo_collection + + self.output_dir = Path(output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + + # Connect to MongoDB + self.client = MongoClient(self.mongo_uri) + self.db = self.client[self.mongo_db] + self.collection = self.db[self.mongo_collection] + + logger.info(f"Connected to MongoDB: {mongo_uri}") + logger.info(f"Output directory: {self.output_dir}") + + # Counter for skipped experiments + self.skipped_no_sampleid = 0 + + # Data containers - consolidated structure + # Main experiments table (merged 1:1 data) + self.experiments = [] + + # 1:N tables (must stay separate) + self.experiment_elements = [] + self.powder_doses = [] + self.temperature_logs = [] + self.xrd_data_points = [] + self.workflow_tasks = [] + + def _build_filter_query(self, experiment_filter: Dict) -> Dict: + """Build MongoDB query from filter criteria""" + query = {} + + if 'types' in experiment_filter and experiment_filter['types']: + # Match experiments starting with any of the prefixes + patterns = [f"^{t}_" for t in experiment_filter['types']] + query["name"] = {"$regex": "|".join(patterns)} + + if 'status' in experiment_filter and experiment_filter['status']: + query["status"] = {"$in": experiment_filter['status']} + + if 'has_xrd' in experiment_filter: + if experiment_filter['has_xrd']: + query["metadata.diffraction_results.sampleid_in_aeris"] = {"$exists": True, "$ne": None} + else: + query["$or"] = [ + {"metadata.diffraction_results.sampleid_in_aeris": {"$exists": False}}, + {"metadata.diffraction_results.sampleid_in_aeris": None} + ] + + if 'date_range' in experiment_filter: + date_query = {} + if 'start' in experiment_filter['date_range']: + date_query["$gte"] = experiment_filter['date_range']['start'] + if 'end' in experiment_filter['date_range']: + date_query["$lte"] = experiment_filter['date_range']['end'] + if date_query: + query["last_updated"] = date_query + + if 'experiment_names' in experiment_filter and experiment_filter['experiment_names']: + query["name"] = {"$in": experiment_filter['experiment_names']} + + return query + + def transform_all(self, limit: Optional[int] = None, + skip_temp_logs: bool = False, + skip_xrd_points: bool = False, + skip_workflow_tasks: bool = False, + experiment_filter: Optional[Dict] = None): + """Transform all experiments to dataframes + + Args: + limit: Maximum number of experiments to process + skip_temp_logs: Skip temperature log data (large arrays) + skip_xrd_points: Skip XRD data points (very large arrays) + skip_workflow_tasks: Skip workflow task history + experiment_filter: Filter criteria for selecting experiments + - types: List of experiment type prefixes (e.g., ['NSC', 'Na']) + - status: List of statuses (e.g., ['completed']) + - has_xrd: Boolean, whether XRD data is required + - date_range: Dict with 'start' and/or 'end' dates + """ + + query = self._build_filter_query(experiment_filter) if experiment_filter else {} + total = self.collection.count_documents(query) + + if limit: + total = min(limit, total) + + logger.info(f"Processing {total} experiments...") + if experiment_filter: + logger.info(f"Applied filter: {experiment_filter}") + + cursor = self.collection.find(query).limit(limit) if limit else self.collection.find(query) + + processed = 0 + for doc in tqdm(cursor, total=total, desc="Transforming experiments"): + try: + if self._transform_experiment(doc, skip_temp_logs, skip_xrd_points, skip_workflow_tasks): + processed += 1 + except Exception as e: + logger.error(f"Error transforming {doc.get('name', 'unknown')}: {e}") + continue + + # Log filtering statistics + logger.info("-" * 40) + logger.info("Filtering Summary:") + logger.info(f" Total experiments in MongoDB: {total}") + logger.info(f" Skipped (no sampleid_in_aeris): {self.skipped_no_sampleid}") + logger.info(f" Processed: {processed}") + logger.info("-" * 40) + logger.info("Transformation complete. Creating dataframes...") + self._save_all_parquet_files() + + def _transform_experiment(self, doc: Dict, skip_temp_logs: bool, + skip_xrd_points: bool, skip_workflow_tasks: bool) -> bool: + """Transform a single experiment document into consolidated format + + Returns: + bool: True if experiment was processed, False if skipped + """ + + exp_name = doc.get('name', '') + metadata = doc.get('metadata', {}) + + # Filter: Skip experiments without sampleid_in_aeris (indicates no XRD performed) + xrd = metadata.get('diffraction_results', {}) + sampleid = xrd.get('sampleid_in_aeris') if xrd else None + if not sampleid: + logger.debug(f"Skipping {exp_name}: No sampleid_in_aeris (no XRD)") + self.skipped_no_sampleid += 1 + return False + + exp_id = str(doc['_id']) + tasks = doc.get('tasks', []) or [] + + # Extract hierarchical type from name: NSC_249_001 -> root=NSC, subgroup=NSC_249 + parts = exp_name.split('_') if '_' in exp_name else [exp_name] + exp_type = parts[0] if parts else 'Unknown' + exp_subgroup = '_'.join(parts[:2]) if len(parts) >= 2 else None + + # ========================================== + # Build consolidated experiment record + # Merging: experiments + heating + recovery + xrd + finalization + dosing + # ========================================== + + experiment_record = { + # Core experiment fields + 'experiment_id': exp_id, + 'name': exp_name, + 'experiment_type': exp_type, + 'experiment_subgroup': exp_subgroup, + 'target_formula': metadata.get('target', ''), + 'last_updated': doc.get('last_updated'), + 'status': self._get_experiment_status(tasks), + 'notes': None, + } + + # --- Heating session fields (prefix: heating_) --- + heating = metadata.get('heating_results', {}) + heating_method = self._determine_heating_method(tasks, heating) + + experiment_record.update({ + 'heating_method': heating_method, + 'heating_temperature': heating.get('heating_temperature'), + 'heating_time': heating.get('heating_time'), + 'heating_cooling_rate': heating.get('cooling_rate'), + 'heating_atmosphere': heating.get('atmosphere'), + 'heating_flow_rate_ml_min': heating.get('flow_rate'), + 'heating_low_temp_calcination': heating.get('low_temperature_calcination', False), + }) + + # --- Powder recovery fields (prefix: recovery_) --- + recovery = metadata.get('recoverpowder_results', {}) + dosing = metadata.get('powderdosing_results', {}) + + # Calculate total dosed mass + total_dosed_mg = self._calculate_total_dosed_mass(dosing) + collected_weight = recovery.get('weight_collected') + yield_pct = None + if total_dosed_mg and total_dosed_mg > 0 and collected_weight: + yield_pct = (collected_weight / total_dosed_mg) * 100 + + experiment_record.update({ + 'recovery_total_dosed_mass_mg': total_dosed_mg if total_dosed_mg > 0 else None, + 'recovery_weight_collected_mg': collected_weight, + 'recovery_yield_percent': yield_pct, + 'recovery_initial_crucible_weight_mg': recovery.get('initial_crucible_weight'), + 'recovery_failure_classification': recovery.get('failure_classification'), + }) + + # --- XRD measurement fields (prefix: xrd_) --- + experiment_record.update({ + 'xrd_sampleid_in_aeris': xrd.get('sampleid_in_aeris'), + 'xrd_holder_index': xrd.get('xrd_holder_index'), + 'xrd_total_mass_dispensed_mg': xrd.get('total_mass_dispensed_mg'), + 'xrd_met_target_mass': xrd.get('met_target_mass'), + }) + + # --- Sample finalization fields (prefix: finalization_) --- + ending = metadata.get('ending_results', {}) + experiment_record.update({ + 'finalization_decoded_sample_id': ending.get('decoded_sample_id'), + 'finalization_successful_labeling': ending.get('successful_labeling'), + 'finalization_storage_location': self._get_storage_location(tasks), + }) + + # --- Dosing session fields (prefix: dosing_) --- + if dosing: + experiment_record.update({ + 'dosing_crucible_position': dosing.get('CruciblePosition'), + 'dosing_crucible_sub_rack': dosing.get('CrucibleSubRack'), + 'dosing_mixing_pot_position': dosing.get('MixingPotPosition'), + 'dosing_ethanol_dispense_volume': dosing.get('EthanolDispenseVolume'), + 'dosing_target_transfer_volume': dosing.get('TargetTransferVolume'), + 'dosing_actual_transfer_mass': dosing.get('ActualTransferMass'), + 'dosing_dac_duration': dosing.get('DACDuration'), + 'dosing_dac_speed': dosing.get('DACSpeed'), + 'dosing_actual_heat_duration': dosing.get('ActualHeatDuration'), + 'dosing_end_reason': dosing.get('EndReason'), + }) + + self.experiments.append(experiment_record) + + # ========================================== + # 1:N Tables (must stay separate) + # ========================================== + + # Experiment elements + elements = metadata.get('elements_present', []) or [] + for element in elements: + self.experiment_elements.append({ + 'experiment_id': exp_id, + 'element_symbol': element, + 'target_atomic_percent': None + }) + + # Powder doses (nested structure) + powders = dosing.get('Powders', []) or [] + for powder in powders: + powder_name = powder.get('PowderName', '') + target_mass = powder.get('TargetMass', 0) + + doses = powder.get('Doses', []) or [] + for idx, dose in enumerate(doses): + actual_mass = dose.get('Mass', 0) + accuracy = ((actual_mass / target_mass * 100) + if target_mass > 0 else None) + + self.powder_doses.append({ + 'experiment_id': exp_id, + 'powder_name': powder_name, + 'target_mass': target_mass, + 'actual_mass': actual_mass, + 'accuracy_percent': accuracy, + 'dose_sequence': idx, + 'head_position': dose.get('HeadPosition'), + 'dose_timestamp': dose.get('TimeStamp') + }) + + # Temperature logs (large arrays, optional) + if not skip_temp_logs and heating: + temp_log = heating.get('temperature_log', {}) or {} + times = temp_log.get('time_minutes', []) or [] + temps = temp_log.get('temperature_celsius', []) or [] + + for idx, (time_min, temp_c) in enumerate(zip(times, temps)): + self.temperature_logs.append({ + 'experiment_id': exp_id, + 'sequence_number': idx, + 'time_minutes': time_min, + 'temperature_celsius': temp_c + }) + + # XRD data points (HUGE arrays, optional) + if not skip_xrd_points and xrd: + twotheta = xrd.get('twotheta', []) or [] + counts = xrd.get('counts', []) or [] + + for idx, (theta, count) in enumerate(zip(twotheta, counts)): + self.xrd_data_points.append({ + 'experiment_id': exp_id, + 'point_index': idx, + 'twotheta': theta, + 'counts': count + }) + + # Workflow tasks (optional) + if not skip_workflow_tasks: + for task in tasks: + self.workflow_tasks.append({ + 'experiment_id': exp_id, + 'task_id': str(task.get('task_id', '')), + 'task_type': task.get('type', ''), + 'status': task.get('status', ''), + 'created_at': task.get('created_at'), + 'started_at': task.get('started_at'), + 'completed_at': task.get('completed_at') + }) + + return True + + def _determine_heating_method(self, tasks: List[Dict], heating: Dict) -> str: + """Determine heating method from task types""" + heating_method = 'none' + for task in tasks: + task_type = task.get('type', '') + if task_type == 'HeatingWithAtmosphere': + heating_method = 'atmosphere' + params = task.get('parameters', {}) + if params.get('atmosphere'): + heating['atmosphere'] = params.get('atmosphere') + heating['flow_rate'] = params.get('flow_rate') + break + elif task_type == 'ManualHeating': + heating_method = 'manual' + break + elif task_type == 'Heating': + heating_method = 'standard' + break + return heating_method + + def _calculate_total_dosed_mass(self, dosing: Dict) -> float: + """Calculate total dosed mass from powder doses (grams -> mg)""" + total_dosed_mg = 0 + powders = dosing.get('Powders', []) or [] + for powder in powders: + doses = powder.get('Doses', []) or [] + for dose in doses: + total_dosed_mg += (dose.get('Mass', 0) or 0) * 1000 # g to mg + return total_dosed_mg + + def _get_experiment_status(self, tasks: List[Dict]) -> str: + """Determine experiment status from tasks""" + if not tasks or tasks is None: + return 'unknown' + + statuses = [t.get('status', '') for t in tasks] + + if all(s == 'COMPLETED' for s in statuses): + return 'completed' + elif any(s == 'ERROR' for s in statuses): + return 'error' + elif any(s in ['RUNNING', 'WAITING'] for s in statuses): + return 'active' + else: + return 'unknown' + + def _get_storage_location(self, tasks: List[Dict]) -> Optional[str]: + """Extract final storage location from Ending task""" + if not tasks: + return None + + for task in tasks: + if task.get('type') == 'Ending': + subtasks = task.get('subtasks', []) + if subtasks: + last_subtask = subtasks[-1] + params = last_subtask.get('parameters', {}) + return params.get('destination') + return None + + def _save_all_parquet_files(self): + """Convert all lists to dataframes and save as parquet""" + + datasets = { + 'experiments': self.experiments, + 'experiment_elements': self.experiment_elements, + 'powder_doses': self.powder_doses, + 'temperature_logs': self.temperature_logs, + 'xrd_data_points': self.xrd_data_points, + 'workflow_tasks': self.workflow_tasks, + } + + for name, data in datasets.items(): + if not data: + logger.warning(f"No data for {name}, skipping...") + continue + + df = pd.DataFrame(data) + output_path = self.output_dir / f"{name}.parquet" + + # Optimize data types for better compression + df = self._optimize_dtypes(df) + + # Save with compression (from config) + df.to_parquet( + output_path, + engine=config.parquet_engine, + compression=config.parquet_compression, + index=False + ) + + file_size_mb = output_path.stat().st_size / (1024 * 1024) + logger.info(f"✓ Saved {name}: {len(df):,} rows, {len(df.columns)} cols, {file_size_mb:.2f} MB") + + # Save metadata + self._save_metadata(datasets) + + def _optimize_dtypes(self, df: pd.DataFrame) -> pd.DataFrame: + """Optimize dataframe dtypes for smaller file size""" + + for col in df.columns: + col_type = df[col].dtype + + # Downcast integers + if col_type == 'int64': + df[col] = pd.to_numeric(df[col], downcast='integer') + + # Downcast floats + elif col_type == 'float64': + df[col] = pd.to_numeric(df[col], downcast='float') + + # Convert to category if low cardinality + elif col_type == 'object': + num_unique = df[col].nunique() + num_total = len(df) + if num_unique / num_total < 0.5: # Less than 50% unique + df[col] = df[col].astype('category') + + return df + + def _save_metadata(self, datasets: Dict[str, List]): + """Save migration metadata""" + + # Get column info for experiments table + exp_columns = list(pd.DataFrame(datasets['experiments']).columns) if datasets['experiments'] else [] + + metadata = { + 'migration_date': datetime.now().isoformat(), + 'version': '2.0', # Consolidated schema version + 'source': { + 'type': 'mongodb', + 'uri': self.mongo_uri, + 'database': self.mongo_db, + 'collection': self.mongo_collection + }, + 'output_directory': str(self.output_dir), + 'schema': { + 'description': 'Consolidated schema - all 1:1 tables merged into experiments.parquet', + 'experiments_columns': len(exp_columns), + 'merged_tables': ['experiments', 'heating_sessions', 'powder_recovery', + 'xrd_measurements', 'sample_finalization', 'dosing_sessions'] + }, + 'datasets': { + name: { + 'rows': len(data), + 'file': f"{name}.parquet" + } + for name, data in datasets.items() if data + } + } + + import json + metadata_path = self.output_dir / 'metadata.json' + with open(metadata_path, 'w') as f: + json.dump(metadata, f, indent=2) + + logger.info(f"✓ Saved metadata: {metadata_path}") + + def close(self): + """Close MongoDB connection""" + self.client.close() + + +def main(): + parser = argparse.ArgumentParser( + description='Transform MongoDB experiments to consolidated Parquet files' + ) + parser.add_argument( + '--output-dir', '-o', + type=Path, + default=OUTPUT_DIR, + help='Output directory for parquet files' + ) + parser.add_argument( + '--mongo-uri', '-m', + default=None, + help=f'MongoDB connection URI (default: from config, currently {config.mongo_uri})' + ) + parser.add_argument( + '--limit', '-l', + type=int, + help='Limit number of experiments to process (for testing)' + ) + parser.add_argument( + '--skip-temp-logs', + action='store_true', + help='Skip temperature log data (saves ~500k rows)' + ) + parser.add_argument( + '--skip-xrd-points', + action='store_true', + help='Skip XRD data points (saves ~4.7M rows)' + ) + parser.add_argument( + '--skip-workflow-tasks', + action='store_true', + help='Skip workflow task history (saves ~3k rows)' + ) + + args = parser.parse_args() + + logger.info("="*60) + logger.info("MongoDB to Parquet Transformation Pipeline (v2 - Consolidated)") + logger.info("="*60) + + try: + transformer = MongoToParquetTransformer( + mongo_uri=args.mongo_uri, + output_dir=args.output_dir + ) + + transformer.transform_all( + limit=args.limit, + skip_temp_logs=args.skip_temp_logs, + skip_xrd_points=args.skip_xrd_points, + skip_workflow_tasks=args.skip_workflow_tasks + ) + + logger.info("="*60) + logger.info("✓ Transformation complete!") + logger.info(f"✓ Parquet files saved to: {args.output_dir}") + logger.info("="*60) + + except Exception as e: + logger.error(f"Pipeline failed: {e}", exc_info=True) + sys.exit(1) + finally: + transformer.close() + + +if __name__ == '__main__': + main() diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/PIPELINE_ARCHITECTURE.md b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/PIPELINE_ARCHITECTURE.md new file mode 100644 index 000000000..024b58a13 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/PIPELINE_ARCHITECTURE.md @@ -0,0 +1,472 @@ +# A-Lab Pipeline Architecture v2 + +## Overview + +The A-Lab pipeline is a product-based data processing system that: + +1. Filters experiments by configurable criteria +2. Transforms MongoDB data to Parquet (uses `mongodb_to_parquet.py` primitive) +3. Runs **auto-discovered** analyses (XRD, powder statistics, custom) +4. Validates data with **auto-discovered** Pydantic schemas +5. Generates schema diagrams per product +6. Uploads parquet files directly to S3 OpenData + +### Key Features + +- **Auto-Discovery**: Schemas and analyses are discovered automatically from their directories +- **Extensible**: Add new data types or analyses by creating files, no code changes needed +- **Configuration-Driven**: Defaults and filters defined in YAML files + +## Quick Start + +### Create a New Data Product + +```bash +# Interactive product creation (uses default schema automatically) +python data/pipeline/run_pipeline.py create + +# Follow prompts to: +# - Name your product (e.g., "reaction_genome") +# - Select experiment groups hierarchically: +# • Root types (NSC, Na, PG, MINES, TRI) - selects ALL experiments +# • Subgroups (NSC_249, Na_123_A) - selects ONLY those experiments +# - Configure filters (has_xrd, status, date_range) +# - Choose analyses to run +# - Schema is loaded from data/products/default_schema.yaml (source of truth) +# - Optionally customize by removing unwanted fields +``` + +### Run the Pipeline + +```bash +# Dry run (default) +python data/pipeline/run_pipeline.py run --product reaction_genome + +# Live upload to S3 OpenData +python data/pipeline/run_pipeline.py run --product reaction_genome --upload + +# Run specific stages only +python data/pipeline/run_pipeline.py run --product reaction_genome --stages filter transform +``` + +### List Products + +```bash +python data/pipeline/run_pipeline.py list +``` + +### Regenerate Pydantic Schema + +```bash +# Regenerate schema.py from config.yaml (auto-syncs on every run) +python data/pipeline/run_pipeline.py regenerate --product reaction_genome +``` + +## Architecture + +``` +┌─────────────────────────────────────────────────────────────────┐ +│ Product Configuration │ +│ data/products/reaction_genome/config.yaml │ +│ │ +│ • Experiment filters (types, status, has_xrd) │ +│ • Product metadata (title, authors, description) │ +│ • Analysis pipeline (xrd_dara, powder_stats, etc.) │ +│ • Schema with units (heatingTemperature: degC) │ +└─────────────────────────────────────────────────────────────────┘ + ↓ +┌─────────────────────────────────────────────────────────────────┐ +│ Filter Experiments │ +│ Hierarchical Selection from MongoDB │ +│ │ +│ Discover: NSC (218) → NSC_249 (12), NSC_250 (8), ... │ +│ Na (215) → Na_123_A (5), Na_123_B (7), ... │ +│ Select: Root types OR specific subgroups │ +│ Apply filters: status=completed, has_xrd=true │ +│ Result: Filtered experiments → experiments.txt │ +└─────────────────────────────────────────────────────────────────┘ + ↓ +┌─────────────────────────────────────────────────────────────────┐ +│ Transform to Parquet │ +│ Create Product-Specific Filtered Data │ +│ │ +│ Each product maintains its own filtered parquet dataset: │ +│ • Data consistency with submitted version │ +│ • No cross-product contamination │ +│ • Schema version tracking per product │ +│ │ +│ Input: Filtered experiment list (from experiments.txt) │ +│ Output: products/{name}/parquet_data/*.parquet (filtered) │ +│ Schema: default_schema.yaml (new) OR config.yaml (submitted) │ +│ Tables: experiments, heating_sessions, xrd_measurements, etc. │ +└─────────────────────────────────────────────────────────────────┘ + ↓ +┌─────────────────────────────────────────────────────────────────┐ +│ Run Analysis Pipeline │ +│ Pluggable Analyzers │ +│ │ +│ • XRDAnalyzer → xrd_success, rwp, num_phases │ +│ • PowderStatisticsAnalyzer → avg_accuracy, total_doses │ +│ • Custom analyzers can be added │ +│ Output: products/reaction_genome/analysis_results/*.parquet │ +└─────────────────────────────────────────────────────────────────┘ + ↓ +┌─────────────────────────────────────────────────────────────────┐ +│ Validate with Pydantic │ +│ Auto-generated Pydantic + Schema Validation │ +│ │ +│ Pydantic-Based Validation Strategy: │ +│ • SchemaManager discovers all YAML schemas in schema/ │ +│ • Generates Pydantic model for each table (*_schema.py) │ +│ • SchemaValidator validates each parquet table row-by-row │ +│ • Leverages Pydantic's type checking and error messages │ +│ │ +│ 1. Auto-discover schemas: experiments, powder_doses, etc. │ +│ 2. Generate Pydantic models from YAML schemas │ +│ 3. Import Pydantic models dynamically at validation time │ +│ 4. Validate every row in every parquet table │ +│ 5. Report validation statistics and detailed errors │ +│ 6. Schema-agnostic: add new table → auto validates │ +└─────────────────────────────────────────────────────────────────┘ + ↓ +┌─────────────────────────────────────────────────────────────────┐ +│ Upload to S3 OpenData │ +│ │ +│ 1. Upload parquet files directly to S3 bucket │ +│ 2. Embed metadata in parquet schema │ +│ 3. Exclude very large files (temperature_logs, xrd_data_points) │ +│ 4. Available via S3 API and HTTP │ +└─────────────────────────────────────────────────────────────────┘ +``` + +## File Structure + +``` +A-Lab_Samples/ +├── data/ +│ ├── config/ # ⭐ Global configuration +│ │ ├── defaults.yaml # Pipeline defaults +│ │ ├── filters.yaml # Filter presets & experiment types +│ │ └── analyses.yaml # Analysis documentation +│ │ +│ ├── products/ # Data product definitions +│ │ ├── base_product.py # Product configuration system +│ │ ├── schema_manager.py # Auto-discovers Pydantic schemas +│ │ ├── schema_validator.py # Validates parquet vs schema +│ │ ├── schema/ # ⭐ AUTO-DISCOVERED schemas +│ │ │ ├── base.py # Shared utilities (ExcludeFromUpload) +│ │ │ ├── experiments.py # Main experiment schema +│ │ │ ├── powder_doses.py # Powder dosing schema +│ │ │ ├── xrd_*.py # XRD-related schemas +│ │ │ └── [your_schema.py] # ⭐ Add your own - auto-discovered! +│ │ │ +│ │ └── {product_name}/ # Product-specific directory +│ │ ├── config.yaml # Product configuration +│ │ ├── experiments.txt # Filtered experiment list +│ │ ├── parquet_data/ # Transformed data +│ │ ├── analysis_results/ # Analysis outputs +│ │ └── SCHEMA_DIAGRAM.md # Auto-generated diagram +│ │ +│ ├── pipeline/ +│ │ ├── run_pipeline.py # Main CLI (create, run, list) +│ │ ├── product_pipeline.py # Pipeline orchestration +│ │ ├── s3_uploader.py # S3 OpenData uploader +│ │ ├── archive/ # Deprecated API-based uploaders +│ │ └── pipeline_state.py # State tracking +│ │ +│ ├── analyses/ # ⭐ AUTO-DISCOVERED analysis plugins +│ │ ├── base_analyzer.py # Base class + built-in analyzers +│ │ └── [your_analyzer.py] # ⭐ Add your own - auto-discovered! +│ │ +│ ├── mongodb_to_parquet.py # Shared primitive: MongoDB → Parquet +│ │ +│ ├── xrd_creation/ # XRD analysis (DARA) +│ │ └── analyze_batch.py # Entry point: run_analysis.sh +│ │ +│ └── tools/ +│ ├── generate_diagram.py # Schema diagram generator +│ └── analyze_mongodb.py # Discover experiment types +│ +├── update_data.sh # MongoDB → Parquet (all data) +├── run_analysis.sh # XRD analysis runner +├── run_dashboard.sh # Dashboard (operates on parquet) +└── run_product_pipeline.sh # Product pipeline CLI wrapper +``` + +### Auto-Discovery Directories + +| Directory | What's Discovered | How to Extend | +| -------------------------- | ----------------- | ------------------------------------- | +| `data/products/schema/` | Pydantic schemas | Add `*.py` with BaseModel class | +| `data/analyses/` | Analysis plugins | Add `*.py` with BaseAnalyzer subclass | +| `data/config/filters.yaml` | Filter presets | Add YAML entry | + +## Schema Management (Auto-Discovered) + +Schemas are **auto-discovered** from `data/products/schema/`. Python-first approach! + +### Schema Directory Structure + +``` +data/products/schema/ +├── base.py # Shared utilities (ExcludeFromUpload) +├── experiments.py # Main experiment schema (~40 fields) +├── experiment_elements.py # Elements per experiment +├── powder_doses.py # Powder dosing data +├── temperature_logs.py # Temperature time series +├── xrd_data_points.py # XRD patterns +├── xrd_refinements.py # DARA analysis results +├── xrd_phases.py # Identified phases +└── [your_schema.py] # ⭐ Auto-discovered! +``` + +### Adding a New Schema + +```python +# data/products/schema/sem_data.py +"""SEM measurement schema""" + +from pydantic import BaseModel, Field + +class SEMData(BaseModel, extra="forbid"): + """SEM measurement data""" + + # Optional: specify table name (defaults to filename) + __schema_table__ = "sem_data" + + experiment_id: str = Field(description="Experiment ID") + image_count: int = Field(description="Number of SEM images") + morphology: str | None = Field(default=None, description="Classified morphology") +``` + +### List Available Schemas + +```bash +# Show all auto-discovered schemas +python data/products/schema_manager.py +``` + +### Marking Fields as Embargoed + +Use `ExcludeFromUpload` to prevent sensitive fields from being uploaded: + +```python +from .base import ExcludeFromUpload + +class MySchema(BaseModel): + public_field: str = Field(description="This is uploaded") + + # This field is NOT uploaded to S3 + sensitive_field: float | None = ExcludeFromUpload( + description="Embargoed until publication" + ) +``` + +### Schema Workflow + +``` +data/products/schema/*.py (Python-first, auto-discovered) + ↓ +SchemaManager loads all schemas automatically + ↓ +SchemaValidator validates parquet data + ↓ +Embargoed fields excluded from upload +``` + +⚠️ **Schemas are Python code** - edit directly, no YAML regeneration needed! + +```yaml +name: reaction_genome +version: 1.0 + +experiment_filter: + types: [NSC, Na] # Experiment prefixes + status: [completed] # Workflow status + has_xrd: true # Must have XRD data + +metadata: + title: 'A-Lab NASICON Synthesis Dataset' + authors: 'A-Lab Team, LBNL' + description: 'High-throughput synthesis data' + references: + - label: github + url: https://github.com/... + +analyses: + - name: xrd_dara + enabled: true + config: + wmin: 10 + wmax: 80 + + - name: powder_statistics + enabled: true + +schema: + heatingTemperature: + unit: degC + description: 'Synthesis temperature' + type: float + required: true + min: 0 + max: 2000 + + recoveryYield: + unit: '%' + description: 'Powder recovery percentage' + type: float + required: false + min: 0 + max: 100 +``` + +## Analysis Plugins (Auto-Discovered) + +Analyses are **auto-discovered** from `data/analyses/`. No registration needed! + +### Creating a Custom Analyzer + +```python +# data/analyses/my_analyzer.py + +from pathlib import Path +import pandas as pd +from base_analyzer import BaseAnalyzer + +class MyCustomAnalyzer(BaseAnalyzer): + """My custom analysis for A-Lab experiments""" + + # Class attributes for discovery (all optional except 'name') + name = "my_analysis" # Required: unique identifier + description = "Calculate my metrics" # Optional: shown in list + cli_flag = "--my-analysis" # Optional: for CLI usage + + def analyze(self, experiments_df: pd.DataFrame, parquet_dir: Path) -> pd.DataFrame: + """Run analysis on experiments""" + results = [] + + for _, exp in experiments_df.iterrows(): + results.append({ + 'experiment_name': exp['name'], + 'my_metric': self._calculate_metric(exp) + }) + + return pd.DataFrame(results) + + def get_output_schema(self): + return { + 'my_metric': { + 'type': 'float', + 'required': True, + 'description': 'My custom metric' + } + } + + def _calculate_metric(self, exp): + # Your analysis logic here + return 0.0 +``` + +### List Available Analyzers + +```bash +# Show all auto-discovered analyzers +python data/analyses/base_analyzer.py +``` + +### Enable in Product Config + +```yaml +# data/products/my_product/config.yaml +analyses: + - name: my_analysis + enabled: true + config: + param1: value1 +``` + +## S3 OpenData Integration + +### AWS Credentials Setup + +Create `.env` file at project root: + +``` +aws_access_key_id=your_key_here +aws_secret_access_key=your_secret_here +``` + +### Data Format + +Parquet files with embedded metadata: + +1. **Column metadata**: Units and descriptions in arrow schema +2. **Project metadata**: Embedded as schema metadata +3. **Efficient storage**: Columnar format for fast partial reads + +### Upload Process + +1. **Prepare parquet files**: Product-specific filtered data +2. **Embed metadata**: Add project info to schema +3. **Upload to S3**: Direct upload to `s3://materialsproject-contribs/alab_synthesis/{product}/` +4. **Exclude large files**: Temperature logs and XRD data points skipped by default + +## State Management + +Pipeline state tracked in `pipeline_runs.parquet`: + +- Run ID and timestamp +- Experiments processed +- Upload status +- Dry run flag + +View state: + +```python +from pipeline_state import PipelineStateManager +mgr = PipelineStateManager() +mgr.get_status_summary() +``` + +## Troubleshooting + +### "No experiments match filter" + +- Check filter criteria in config.yaml +- Use `python data/tools/analyze_mongodb.py` to explore available types +- Check `data/config/filters.yaml` for preset options + +### "Validation errors" + +- Check Pydantic schema matches your data +- Run `python data/products/schema_manager.py` to see loaded schemas +- Verify min/max ranges are appropriate + +### "S3 upload failed" + +- Verify AWS credentials in .env +- Check AWS access permissions +- Ensure boto3 is installed: `pip install boto3` +- Verify network connectivity to S3 + +### "Analysis not found" + +- Run `python data/analyses/base_analyzer.py` to see available analyzers +- Ensure analyzer file is in `data/analyses/` directory +- Check that analyzer class inherits from `BaseAnalyzer` +- Verify `name` class attribute is set + +### "Schema not found" + +- Run `python data/products/schema_manager.py` to see discovered schemas +- Ensure schema file is in `data/products/schema/` directory +- Check that schema class inherits from `pydantic.BaseModel` +- Verify file is not in SKIP_FILES list (base.py, **init**.py) + +## Best Practices + +1. **Start with dry runs**: Test pipeline before uploading +2. **Validate early**: Use Pydantic to catch issues +3. **Modular analyses**: Keep analyzers independent +4. **Version configs**: Track changes to product definitions diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/README.md b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/README.md new file mode 100644 index 000000000..e696dd524 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/README.md @@ -0,0 +1,267 @@ +# A-Lab Data Pipeline + +Orchestrates data flow from MongoDB → Parquet → XRD Analysis → S3 OpenData. + +## 📚 Documentation + +- **[PARQUET_STRUCTURE.md](./PARQUET_STRUCTURE.md)** - Complete guide to parquet files (14 tables, relationships, usage) +- **[PIPELINE_ARCHITECTURE.md](./PIPELINE_ARCHITECTURE.md)** - Product-based pipeline design +- **[requirements.txt](./requirements.txt)** - Python dependencies + +## Quick Start + +```bash +# Create a new data product +python run_pipeline.py create + +# List available products +python run_pipeline.py list + +# Run pipeline (dry run) +python run_pipeline.py run --product reaction_genome + +# Run pipeline with live upload to S3 +python run_pipeline.py run --product reaction_genome --upload + +# Check status +python run_pipeline.py status --product reaction_genome +``` + +## Pipeline Phases + +| Phase | Script | Output | +| --------------- | ------------------ | --------------------------------------------------------------------- | +| 1. MongoDB Sync | `update_data.sh` | 14 parquet files (see [PARQUET_STRUCTURE.md](./PARQUET_STRUCTURE.md)) | +| 2. XRD Analysis | `analyze_batch.py` | `xrd_refinements.parquet`, `xrd_phases.parquet` | +| 3. S3 Upload | `s3_uploader.py` | Parquet files on s3://materialsproject-contribs | + +### Parquet Files Generated + +The pipeline creates **14 normalized parquet tables** (not one big file): + +| Category | File | Rows | Purpose | +| ----------------- | --------------------------- | ----- | ------------------------------------ | +| **Core** | experiments.parquet | 592 | Main experiment metadata | +| **Related (1:1)** | heating_sessions.parquet | 592 | Heating parameters | +| | powder_recovery.parquet | 592 | Powder collection data | +| | xrd_measurements.parquet | 592 | XRD measurement metadata | +| | dosing_sessions.parquet | 588 | Dosing session parameters | +| | sample_finalization.parquet | 592 | Labeling & storage | +| **Details (1:N)** | experiment_elements.parquet | 2,666 | Elements per experiment | +| | powder_doses.parquet | 2,245 | Individual dose events | +| | workflow_tasks.parquet | 3,074 | Task execution history | +| | temperature_logs.parquet\* | 501K | Temperature readings | +| | xrd_data_points.parquet\* | 4.7M | Raw XRD patterns | +| **Analysis** | xrd_refinements.parquet | 576 | All XRD results (success + failures) | +| | xrd_phases.parquet | 139 | Identified phases (successes only) | + +\* _Optional: Use `--fast` flag to skip these large arrays_ + +**Why multiple files?** Efficient loading - only read what you need. Most queries load <100 KB instead of 27 MB. + +📖 **See [PARQUET_STRUCTURE.md](./PARQUET_STRUCTURE.md) for complete details on structure, relationships, and usage examples.** + +## Files + +``` +data/pipeline/ +├── run_pipeline.py # Product-based CLI (create, run, list) +├── product_pipeline.py # Pipeline orchestration per product +├── pipeline_state.py # State tracking (Parquet-based) +├── s3_uploader.py # S3 OpenData uploader +├── archive/ # Deprecated MPContribs API files +│ ├── mpcontribs_manager.py +│ └── mpcontribs_uploader.py +├── pipeline_runs.parquet # Run history (auto-generated) +├── pipeline.log # Execution logs +└── requirements.txt # Dependencies +``` + +## Configuration + +### AWS Credentials + +Set in `.env` at project root for S3 uploads: + +``` +aws_access_key_id=your_key_here +aws_secret_access_key=your_secret_here +``` + +Or export as environment variables: + +```bash +export AWS_ACCESS_KEY_ID=your_key_here +export AWS_SECRET_ACCESS_KEY=your_secret_here +``` + +## CLI Usage + +### Product Management + +```bash +# Create a new data product (interactive) +python run_pipeline.py create + +# List all available products +python run_pipeline.py list + +# Show status of a product +python run_pipeline.py status --product reaction_genome + +# Validate product configuration +python run_pipeline.py validate --product reaction_genome + +# Regenerate Pydantic schema from config +python run_pipeline.py regenerate --product reaction_genome +``` + +### Running Pipelines + +```bash +# Dry run - preview what would be uploaded +python run_pipeline.py run --product reaction_genome + +# Live upload to S3 OpenData +python run_pipeline.py run --product reaction_genome --upload + +# Run specific pipeline stages only +python run_pipeline.py run --product reaction_genome --stages filter transform analyze +``` + +### Using the Shell Script + +```bash +# Use the convenience script (activates venv automatically) +./run_product_pipeline.sh create +./run_product_pipeline.sh list +./run_product_pipeline.sh run --product reaction_genome +./run_product_pipeline.sh run --product reaction_genome --upload +``` + +## State Tracking + +Pipeline state is stored in `pipeline_runs.parquet`: + +| Column | Description | +| ----------------- | ------------------------------------------------------ | +| `run_id` | Unique run identifier | +| `run_timestamp` | When the run occurred | +| `phase` | Pipeline phase (mongodb_sync, xrd_analysis, s3_upload) | +| `experiment_name` | Experiment processed | +| `status` | success / failed / pending | +| `dry_run` | Whether this was a dry run | +| `uploaded_to_s3` | Upload status | +| `s3_key` | S3 object key if uploaded | + +### Incremental Processing + +The pipeline automatically tracks: + +- **New experiments**: Not yet processed +- **Updated experiments**: `last_updated` changed since last run +- **Pending uploads**: Analyzed but not uploaded + +## S3 OpenData Integration + +### Parquet Upload + +Parquet files are uploaded directly to AWS S3 OpenData bucket: + +- **Bucket**: `s3://materialsproject-contribs/alab_synthesis/{product_name}/` +- **Format**: Parquet files with embedded metadata in schema +- **Access**: Available via AWS S3 API and HTTP + +### Files Uploaded + +All product parquet files except very large ones: + +- `experiments.parquet` - Main experiment data +- `experiment_elements.parquet` - Element compositions +- `powder_doses.parquet` - Dosing details +- `heating_sessions.parquet` - Heating parameters +- `xrd_refinements.parquet` - XRD analysis results +- `xrd_phases.parquet` - Identified phases +- And more... + +**Excluded** (too large): `temperature_logs.parquet`, `xrd_data_points.parquet` + +### Metadata + +Metadata is embedded in parquet schema following MPContribs conventions: + +```python +# Embedded in arrow table metadata +{ + "project.name": "product_name", + "project.type": "alab_synthesis", + "columns.heating_temperature.unit": "degC", + "columns.heating_temperature.description": "Heating temperature" +} +``` + +## Dashboard Integration + +The dashboard automatically reads pipeline status: + +```python +from parquet_data_loader import ParquetDataLoader + +loader = ParquetDataLoader() +status = loader.get_pipeline_status() +# {'total_experiments': 576, 'analyzed_experiments': 450, 'uploaded_experiments': 200, ...} +``` + +## Troubleshooting + +### "No experiments to upload" + +Make sure your product has experiments configured. Check with: + +```bash +python run_pipeline.py status --product your_product +``` + +### "AWS credentials not set" + +Check `.env` file exists at project root with: + +``` +aws_access_key_id=your_key_here +aws_secret_access_key=your_secret_here +``` + +### "Import error: boto3" + +Install the AWS client: + +```bash +pip install boto3 python-dotenv +``` + +### View Logs + +```bash +# Pipeline log +cat data/pipeline/pipeline.log + +# XRD analysis log +cat data/xrd_creation/batch_analysis.log +``` + +## Dependencies + +All dependencies are consolidated in the main requirements file: + +```bash +pip install -r data/requirements.txt +``` + +Key dependencies for pipeline: + +- `pandas>=2.0` - Data processing +- `pyarrow>=14.0` - Parquet files +- `boto3>=1.40.0` - S3 uploads +- `python-dotenv>=1.0` - Environment variables +- `pydantic>=2.0` - Data validation +- `PyYAML>=6.0` - Configuration files diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/__init__.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/__init__.py new file mode 100644 index 000000000..0d09c9bd4 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/__init__.py @@ -0,0 +1,18 @@ +""" +A-Lab Data Pipeline - Product-Based Pipeline System + +Orchestrates data flow: MongoDB → Parquet → XRD Analysis → S3 OpenData + +Usage: + python run_pipeline.py create # Create new product + python run_pipeline.py list # List products + python run_pipeline.py run --product # Run pipeline (dry run) + python run_pipeline.py run --product --upload # Live upload to S3 + python run_pipeline.py status --product # Show status +""" + +from .pipeline_state import PipelineStateManager +from .s3_uploader import S3Uploader, upload_product_to_s3 + +__all__ = ['PipelineStateManager', 'S3Uploader', 'upload_product_to_s3'] + diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/mpcontribs_setup.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/mpcontribs_setup.py new file mode 100644 index 000000000..9cfc2a2b3 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/mpcontribs_setup.py @@ -0,0 +1,334 @@ +#!/usr/bin/env python3 +""" +MPContribs Project Setup + +Creates/updates empty MPContribs project with metadata only. +Does NOT upload data - that goes to S3. + +This module handles: +- Creating new MPContribs projects +- Updating project metadata (title, authors, description, references) +- No column definitions +- No data upload +""" + +import os +import logging +from typing import Dict, Optional, List +from pathlib import Path + +logger = logging.getLogger(__name__) + + +class MPContribsProjectManager: + """Manage MPContribs project metadata (empty config only)""" + + def __init__(self, api_key: Optional[str] = None): + """ + Initialize MPContribs manager. + + Args: + api_key: Materials Project API key (or set MP_API_KEY env var) + """ + self.api_key = api_key or os.getenv('MP_API_KEY') or os.getenv('ALAB_MP_API_KEY') + + if not self.api_key: + raise ValueError( + "MP API key required. Set MP_API_KEY or ALAB_MP_API_KEY environment variable, " + "or pass api_key parameter" + ) + + self._mpr = None + self._client = None + + def _get_mprester(self): + """Get MPRester instance (lazy load)""" + if self._mpr is None: + try: + from mp_api.client import MPRester + self._mpr = MPRester(api_key=self.api_key) + except ImportError: + raise ImportError( + "mp-api package required. Install with: pip install mp-api>=0.45.13" + ) + return self._mpr + + def _get_client(self, project_name: str): + """Get ContribsClient instance (lazy load)""" + try: + from mpcontribs.client import Client as ContribsClient + return ContribsClient(project=project_name, apikey=self.api_key) + except ImportError: + raise ImportError( + "mpcontribs-client package required. Install with: pip install mpcontribs-client>=5.10.4" + ) + + def project_exists(self, project_name: str) -> bool: + """ + Check if project exists in MPContribs. + + Args: + project_name: Name of the project + + Returns: + True if project exists, False otherwise + """ + try: + mpr = self._get_mprester() + projects = mpr.contribs.get_projects() + return project_name in [p.get('name') for p in projects] + except Exception as e: + logger.error(f"Error checking project existence: {e}") + return False + + def create_project( + self, + name: str, + title: str, + authors: str, + description: str, + references: Optional[List[Dict[str, str]]] = None, + dry_run: bool = False + ) -> bool: + """ + Create empty MPContribs project with metadata. + + Args: + name: Project name (identifier) + title: Project title + authors: Comma-separated author names + description: Project description + references: Optional list of references [{"label": "...", "url": "..."}] + dry_run: If True, only simulate the operation + + Returns: + True if successful, False otherwise + """ + if dry_run: + logger.info(f"[DRY RUN] Would create MPContribs project: {name}") + logger.info(f" Title: {title}") + logger.info(f" Authors: {authors}") + logger.info(f" Description: {description[:100]}...") + if references: + logger.info(f" References: {len(references)} links") + return True + + try: + mpr = self._get_mprester() + + # Check if project already exists + if self.project_exists(name): + logger.info(f"Project '{name}' already exists, updating instead") + return self.update_project(name, title, authors, description, references) + + # Create project + logger.info(f"Creating MPContribs project: {name}") + mpr.contribs.create_project( + name=name, + title=title, + authors=authors, + description=description + ) + + # Add references if provided + if references: + client = self._get_client(name) + client.update_project({"references": references}) + logger.info(f" Added {len(references)} references") + + logger.info(f"✓ Successfully created project: {name}") + logger.info(f" View at: https://next-gen.materialsproject.org/contribs/projects/{name}") + return True + + except Exception as e: + logger.error(f"Error creating project: {e}") + return False + + def update_project( + self, + name: str, + title: Optional[str] = None, + authors: Optional[str] = None, + description: Optional[str] = None, + references: Optional[List[Dict[str, str]]] = None, + dry_run: bool = False + ) -> bool: + """ + Update existing MPContribs project metadata. + + Args: + name: Project name + title: New title (optional) + authors: New authors (optional) + description: New description (optional) + references: New references (optional) + dry_run: If True, only simulate the operation + + Returns: + True if successful, False otherwise + """ + if dry_run: + logger.info(f"[DRY RUN] Would update MPContribs project: {name}") + if title: + logger.info(f" Title: {title}") + if authors: + logger.info(f" Authors: {authors}") + if description: + logger.info(f" Description: {description[:100]}...") + if references: + logger.info(f" References: {len(references)} links") + return True + + try: + if not self.project_exists(name): + logger.error(f"Project '{name}' does not exist. Create it first.") + return False + + client = self._get_client(name) + + # Build update dict + updates = {} + if title: + updates['title'] = title + if authors: + updates['authors'] = authors + if description: + updates['description'] = description + if references: + updates['references'] = references + + if not updates: + logger.info("No updates to apply") + return True + + logger.info(f"Updating MPContribs project: {name}") + client.update_project(updates) + logger.info(f"✓ Successfully updated project: {name}") + return True + + except Exception as e: + logger.error(f"Error updating project: {e}") + return False + + def delete_project(self, name: str, dry_run: bool = False) -> bool: + """ + Delete MPContribs project. + + Args: + name: Project name + dry_run: If True, only simulate the operation + + Returns: + True if successful, False otherwise + """ + if dry_run: + logger.info(f"[DRY RUN] Would delete MPContribs project: {name}") + return True + + try: + if not self.project_exists(name): + logger.warning(f"Project '{name}' does not exist") + return False + + client = self._get_client(name) + logger.warning(f"Deleting MPContribs project: {name}") + client.delete_project() + logger.info(f"✓ Successfully deleted project: {name}") + return True + + except Exception as e: + logger.error(f"Error deleting project: {e}") + return False + + +def setup_mpcontribs_project( + product_config: Dict, + dry_run: bool = False, + api_key: Optional[str] = None +) -> bool: + """ + Setup MPContribs project from product configuration. + + Args: + product_config: Product configuration dict + dry_run: If True, simulate the operation + api_key: MP API key (optional) + + Returns: + True if successful, False otherwise + """ + manager = MPContribsProjectManager(api_key=api_key) + + # Extract metadata from product config + name = product_config.get('name') + metadata = product_config.get('metadata', {}) + + title = metadata.get('title') or f"A-Lab {name.replace('_', ' ').title()}" + authors = metadata.get('authors') or "A-Lab Team" + description = metadata.get('description') or f"Automated synthesis data for {name}" + + # Build references list + references = [] + if metadata.get('doi'): + references.append({ + 'label': 'doi', + 'url': f"https://doi.org/{metadata['doi']}" + }) + + # Add any additional references from config + if metadata.get('references'): + for ref in metadata['references']: + if isinstance(ref, dict) and 'label' in ref and 'url' in ref: + references.append(ref) + + # Always add S3 data location reference + from config_loader import get_config + config = get_config() + s3_url = f"s3://{config.s3_bucket}/{config.s3_prefix}/{name}/" + references.append({ + 'label': 'data', + 'url': s3_url.replace('s3://', 'https://s3.amazonaws.com/') + }) + + # Create/update project + return manager.create_project( + name=name, + title=title, + authors=authors, + description=description, + references=references if references else None, + dry_run=dry_run + ) + + +if __name__ == '__main__': + # Test script + import sys + + if len(sys.argv) < 2: + print("Usage: python mpcontribs_setup.py [--dry-run]") + sys.exit(1) + + project_name = sys.argv[1] + dry_run = '--dry-run' in sys.argv + + manager = MPContribsProjectManager() + + # Test project creation + success = manager.create_project( + name=project_name, + title=f"Test Project {project_name}", + authors="Test Author", + description="Test description for MPContribs project", + references=[ + {"label": "github", "url": "https://github.com/example/repo"} + ], + dry_run=dry_run + ) + + if success: + print(f"✓ Project setup successful") + else: + print(f"✗ Project setup failed") + sys.exit(1) + diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/pipeline_state.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/pipeline_state.py new file mode 100755 index 000000000..928a5c199 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/pipeline_state.py @@ -0,0 +1,292 @@ +#!/usr/bin/env python3 +""" +Pipeline State Manager - Tracks all pipeline runs in Parquet format + +The pipeline_runs.parquet file serves as an audit log and enables: +- Incremental processing (only process new/updated experiments) +- Upload tracking (know what's been uploaded to MPContribs) +- Run history and debugging +""" + +import uuid +from datetime import datetime +from pathlib import Path +from typing import List, Optional, Dict +import pandas as pd +import numpy as np + +PIPELINE_DIR = Path(__file__).parent +STATE_FILE = PIPELINE_DIR / "pipeline_runs.parquet" +PARQUET_DATA_DIR = PIPELINE_DIR.parent / "parquet" + + +class PipelineStateManager: + """Manage pipeline state in Parquet format""" + + def __init__(self, state_file: Path = STATE_FILE): + self.state_file = state_file + self._runs_df = None + + @property + def runs(self) -> pd.DataFrame: + """Lazy load runs dataframe""" + if self._runs_df is None: + if self.state_file.exists(): + self._runs_df = pd.read_parquet(self.state_file) + else: + self._runs_df = pd.DataFrame(columns=[ + 'run_id', 'run_timestamp', 'run_type', 'phase', + 'experiment_name', 'experiment_last_updated', + 'status', 'error_message', 'duration_seconds', + 'dry_run', 'uploaded_to_mpcontribs', 'mpcontribs_contribution_id' + ]) + return self._runs_df + + def save(self): + """Save state to parquet""" + self.runs.to_parquet(self.state_file, index=False) + + def get_experiments_to_process(self) -> List[str]: + """ + Get experiments that need processing based on: + 1. New experiments (not in state) + 2. Updated experiments (last_updated changed) + """ + experiments_file = PARQUET_DATA_DIR / "experiments.parquet" + if not experiments_file.exists(): + return [] + + experiments_df = pd.read_parquet(experiments_file) + + # Get latest successful run per experiment + if len(self.runs) == 0: + return experiments_df['name'].tolist() + + successful_runs = self.runs[ + (self.runs['phase'] == 'xrd_analysis') & + (self.runs['status'] == 'success') + ] + + if len(successful_runs) == 0: + return experiments_df['name'].tolist() + + latest_per_exp = successful_runs.groupby('experiment_name').agg({ + 'experiment_last_updated': 'max' + }).reset_index() + + # Merge with current experiments + merged = experiments_df.merge( + latest_per_exp, + left_on='name', + right_on='experiment_name', + how='left', + suffixes=('', '_processed') + ) + + # Find new or updated experiments + needs_processing = merged[ + (merged['experiment_last_updated'].isna()) | + (pd.to_datetime(merged['last_updated']) > pd.to_datetime(merged['experiment_last_updated'])) + ] + + return needs_processing['name'].tolist() + + def get_experiments_to_upload(self) -> List[str]: + """Get experiments that have been analyzed but not uploaded to MPContribs""" + # Get all experiments with successful XRD refinements + refinements_file = PARQUET_DATA_DIR / "xrd_refinements.parquet" + if not refinements_file.exists(): + return [] + + refinements_df = pd.read_parquet(refinements_file) + analyzed = set(refinements_df[refinements_df['success'] == True]['experiment_name'].unique()) + + if len(self.runs) == 0: + return list(analyzed) + + # Get uploaded experiments (non-dry-run) + uploaded = set(self.runs[ + (self.runs['phase'] == 'mpcontribs_upload') & + (self.runs['status'] == 'success') & + (self.runs['dry_run'] == False) + ]['experiment_name'].unique()) + + return list(analyzed - uploaded) + + def record_run( + self, + run_type: str, + phases: List[str], + experiments: Optional[List[str]] = None, + dry_run: bool = True + ) -> str: + """ + Record a pipeline run (batch record for all phases/experiments) + + Returns: + run_id for tracking + """ + run_id = str(uuid.uuid4())[:8] + timestamp = datetime.now() + + if experiments is None: + experiments = ['_batch_'] # Placeholder for batch operations + + new_rows = [] + for exp in experiments: + for phase in phases: + new_rows.append({ + 'run_id': run_id, + 'run_timestamp': timestamp, + 'run_type': run_type, + 'phase': phase, + 'experiment_name': exp, + 'experiment_last_updated': timestamp, + 'status': 'pending', + 'error_message': None, + 'duration_seconds': 0.0, + 'dry_run': dry_run, + 'uploaded_to_mpcontribs': False, + 'mpcontribs_contribution_id': None + }) + + if new_rows: + new_df = pd.DataFrame(new_rows) + if len(self.runs) == 0: + self._runs_df = new_df + else: + self._runs_df = pd.concat([self.runs, new_df], ignore_index=True) + self.save() + + return run_id + + def update_experiment_status( + self, + run_id: str, + experiment_name: str, + phase: str, + status: str, + duration_seconds: float = 0.0, + error_message: Optional[str] = None, + mpcontribs_id: Optional[str] = None + ): + """Update status for a specific experiment/phase in a run""" + mask = ( + (self.runs['run_id'] == run_id) & + (self.runs['experiment_name'] == experiment_name) & + (self.runs['phase'] == phase) + ) + + if not mask.any(): + # Add new row if doesn't exist + new_row = { + 'run_id': run_id, + 'run_timestamp': datetime.now(), + 'run_type': 'manual', + 'phase': phase, + 'experiment_name': experiment_name, + 'experiment_last_updated': datetime.now(), + 'status': status, + 'error_message': error_message, + 'duration_seconds': duration_seconds, + 'dry_run': False, + 'uploaded_to_mpcontribs': mpcontribs_id is not None, + 'mpcontribs_contribution_id': mpcontribs_id + } + self._runs_df = pd.concat([self.runs, pd.DataFrame([new_row])], ignore_index=True) + else: + self._runs_df.loc[mask, 'status'] = status + self._runs_df.loc[mask, 'duration_seconds'] = duration_seconds + + if error_message: + self._runs_df.loc[mask, 'error_message'] = error_message + + if mpcontribs_id: + self._runs_df.loc[mask, 'mpcontribs_contribution_id'] = mpcontribs_id + self._runs_df.loc[mask, 'uploaded_to_mpcontribs'] = True + + self.save() + + def get_run_summary(self, run_id: Optional[str] = None) -> Dict: + """Get summary statistics for a run (or latest run if not specified)""" + if len(self.runs) == 0: + return {'error': 'No runs found'} + + if run_id is None: + run_id = self.runs['run_id'].iloc[-1] + + run_data = self.runs[self.runs['run_id'] == run_id] + + if len(run_data) == 0: + return {'error': f'Run {run_id} not found'} + + return { + 'run_id': run_id, + 'timestamp': str(run_data['run_timestamp'].iloc[0]), + 'total_experiments': run_data['experiment_name'].nunique(), + 'phases': run_data['phase'].unique().tolist(), + 'success_count': len(run_data[run_data['status'] == 'success']), + 'failed_count': len(run_data[run_data['status'] == 'failed']), + 'pending_count': len(run_data[run_data['status'] == 'pending']), + 'dry_run': bool(run_data['dry_run'].iloc[0]), + 'total_duration': float(run_data['duration_seconds'].sum()) + } + + def get_last_run_timestamp(self) -> Optional[datetime]: + """Get timestamp of last successful run""" + if len(self.runs) == 0: + return None + + successful = self.runs[self.runs['status'] == 'success'] + if len(successful) == 0: + return None + + return successful['run_timestamp'].max() + + def get_status_summary(self) -> Dict: + """Get overall pipeline status summary""" + experiments_file = PARQUET_DATA_DIR / "experiments.parquet" + refinements_file = PARQUET_DATA_DIR / "xrd_refinements.parquet" + + total_experiments = 0 + analyzed_experiments = 0 + + if experiments_file.exists(): + total_experiments = len(pd.read_parquet(experiments_file)) + + if refinements_file.exists(): + ref_df = pd.read_parquet(refinements_file) + analyzed_experiments = len(ref_df[ref_df['success'] == True]) + + uploaded_experiments = 0 + if len(self.runs) > 0: + uploaded_experiments = len(self.runs[ + (self.runs['phase'] == 'mpcontribs_upload') & + (self.runs['status'] == 'success') & + (self.runs['dry_run'] == False) + ]['experiment_name'].unique()) + + return { + 'total_experiments': total_experiments, + 'analyzed_experiments': analyzed_experiments, + 'uploaded_experiments': uploaded_experiments, + 'pending_analysis': len(self.get_experiments_to_process()), + 'pending_upload': len(self.get_experiments_to_upload()), + 'last_run': str(self.get_last_run_timestamp()) if self.get_last_run_timestamp() else None, + 'total_runs': self.runs['run_id'].nunique() if len(self.runs) > 0 else 0 + } + + +if __name__ == '__main__': + mgr = PipelineStateManager() + + print("Pipeline State Summary") + print("=" * 50) + + status = mgr.get_status_summary() + for key, value in status.items(): + print(f" {key}: {value}") + + print("\nExperiments to process:", len(mgr.get_experiments_to_process())) + print("Experiments to upload:", len(mgr.get_experiments_to_upload())) + diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/product_pipeline.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/product_pipeline.py new file mode 100644 index 000000000..6c44b3058 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/product_pipeline.py @@ -0,0 +1,691 @@ +#!/usr/bin/env python3 +""" +Product Pipeline Runner + +Orchestrates the full pipeline for a data product: +1. Load product configuration +2. Filter experiments based on criteria +3. Transform MongoDB to Parquet (consolidated schema) +4. Run configured analyses +5. Validate with Pydantic +6. Upload parquet files directly to S3 OpenData + +Output Files: +- experiments.parquet: Main table with ALL 1:1 data (~45 columns) +- experiment_elements.parquet: Elements per experiment (1:N) +- powder_doses.parquet: Individual powder doses (1:N) +- temperature_logs.parquet: Temperature readings (1:N, optional) +- xrd_data_points.parquet: Raw XRD patterns (1:N, optional) +- workflow_tasks.parquet: Task execution history (1:N, optional) +""" + +import argparse +import sys +import logging +from pathlib import Path +from datetime import datetime +from typing import Dict, List, Optional +import pandas as pd +import shutil + +# Add parent directories to path +sys.path.insert(0, str(Path(__file__).parent.parent)) +sys.path.insert(0, str(Path(__file__).parent.parent / "products")) +sys.path.insert(0, str(Path(__file__).parent.parent / "analyses")) +sys.path.insert(0, str(Path(__file__).parent.parent / "config")) + +from base_product import ProductConfig, ProductManager +from base_analyzer import AnalysisPluginManager +from mongodb_to_parquet import MongoToParquetTransformer +from pipeline_state import PipelineStateManager +from s3_uploader import S3Uploader, upload_product_to_s3 +from mpcontribs_setup import setup_mpcontribs_project +from config_loader import get_config + +# Add tools directory for diagram generation +sys.path.insert(0, str(Path(__file__).parent.parent / "tools")) + +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(levelname)s - %(message)s' +) +logger = logging.getLogger(__name__) + +PARQUET_DATA_DIR = Path("data/parquet") + + +class ProductPipeline: + """Manages pipeline execution for a data product""" + + def __init__(self, product_name: str): + self.product_name = product_name + self.product_dir = Path(f"data/products/{product_name}") + self.config: Optional[ProductConfig] = None + self.state_manager = PipelineStateManager() + + def load_config(self) -> bool: + """Load product configuration""" + config_file = self.product_dir / "config.yaml" + + if not config_file.exists(): + logger.error(f"Config file not found: {config_file}") + return False + + try: + self.config = ProductConfig.load(config_file) + logger.info(f"Loaded config for product: {self.product_name}") + return True + except Exception as e: + logger.error(f"Failed to load config: {e}") + return False + + def filter_experiments(self) -> pd.DataFrame: + """ + Filter experiments from MongoDB based on configuration + + Returns: + DataFrame with filtered experiments + """ + from pymongo import MongoClient + + # Connect to MongoDB (using config loader) + config = get_config() + client = MongoClient(config.mongo_uri) + db = client[config.mongo_db] + collection = db[config.mongo_collection] + + # Build query from filter + query = self.config.experiment_filter.to_mongo_query() + + # Get matching experiments + experiments = list(collection.find(query, { + '_id': 1, + 'name': 1, + 'metadata.target': 1, + 'last_updated': 1, + 'status': 1 + })) + + client.close() + + # Convert to DataFrame + df = pd.DataFrame(experiments) + df['experiment_id'] = df['_id'].astype(str) + + # Extract target formula from metadata + df['target_formula'] = df.get('metadata', pd.Series()).apply( + lambda x: x.get('target', '') if isinstance(x, dict) else '' + ) + + # Extract experiment_type and experiment_subgroup from name + def extract_type_and_subgroup(name): + parts = name.split('_') + exp_type = parts[0] if parts else name + exp_subgroup = '_'.join(parts[:2]) if len(parts) >= 2 else exp_type + return pd.Series({'experiment_type': exp_type, 'experiment_subgroup': exp_subgroup}) + + df[['experiment_type', 'experiment_subgroup']] = df['name'].apply(extract_type_and_subgroup) + + logger.info(f"Found {len(df)} experiments matching filter") + + # Save list to product directory + experiment_list_file = self.product_dir / "experiments.txt" + df['name'].to_csv(experiment_list_file, index=False, header=False) + logger.info(f"Saved experiment list to {experiment_list_file}") + + return df + + def transform_to_parquet(self, experiments_df: pd.DataFrame) -> bool: + """ + Transform filtered experiments to consolidated Parquet format + + Uses the same transformation as update_data.sh for consistency. + """ + experiment_names = experiments_df['name'].tolist() + product_parquet_dir = self.product_dir / "parquet_data" + + # Check if product parquet already exists with correct experiments + metadata_file = product_parquet_dir / "metadata.json" + if metadata_file.exists(): + import json + with open(metadata_file) as f: + metadata = json.load(f) + existing_count = metadata.get('datasets', {}).get('experiments', {}).get('rows', 0) + if existing_count == len(experiment_names): + logger.info(f"✓ Product parquet already exists ({len(experiment_names)} experiments)") + logger.info(" Skipping transformation - using existing data") + return True + + # Create product-specific filtered parquet + logger.info(f"Creating filtered parquet for {len(experiment_names)} experiments...") + + experiment_filter = {'experiment_names': experiment_names} + + try: + transformer = MongoToParquetTransformer(output_dir=product_parquet_dir) + transformer.transform_all(experiment_filter=experiment_filter) + transformer.close() + + logger.info(f"✓ Created product-specific parquet ({len(experiment_names)} experiments)") + logger.info(f" Location: {product_parquet_dir}") + return True + + except Exception as e: + logger.error(f"Parquet transformation failed: {e}") + return False + + def run_analyses(self, experiments_df: pd.DataFrame) -> pd.DataFrame: + """ + Run configured analyses on parquet data + + Args: + experiments_df: Experiments to analyze + + Returns: + DataFrame with analysis results (if any) + """ + parquet_dir = self.product_dir / "parquet_data" + output_dir = self.product_dir / "analysis_results" + + # Ensure output directory exists + output_dir.mkdir(parents=True, exist_ok=True) + + # If no analyses configured, just return + if not self.config.analyses: + logger.info("No analyses configured") + return pd.DataFrame() + + # Run analyses using plugin manager + plugin_manager = AnalysisPluginManager() + + results_df = plugin_manager.run_analyses( + self.config.analyses, + experiments_df, + parquet_dir, + output_dir + ) + + logger.info(f"✓ Completed {len(self.config.analyses)} analyses") + + # Save analysis results (only if there are results) + if len(results_df) > 0 and len(self.config.analyses) > 0: + results_file = output_dir / "analysis_results.parquet" + results_df.to_parquet(results_file, index=False) + logger.info(f"✓ Saved analysis results: {len(results_df)} rows") + + return results_df + + def load_experiments_data(self) -> pd.DataFrame: + """ + Load experiment data from consolidated parquet file + + Returns: + DataFrame with all experiment data from experiments.parquet + """ + experiments_file = self.product_dir / "parquet_data" / "experiments.parquet" + + if not experiments_file.exists(): + logger.error(f"Experiments file not found: {experiments_file}") + return pd.DataFrame() + + return pd.read_parquet(experiments_file) + + def validate_data(self, data_df: pd.DataFrame) -> bool: + """ + Validate data against Pydantic schema + + Args: + data_df: Data to validate + + Returns: + True if all records valid + """ + # Load Pydantic schema + schema_file = self.product_dir / "schema.py" + + if not schema_file.exists(): + logger.warning("No Pydantic schema found, skipping validation") + return True + + # Import the schema module + import importlib.util + spec = importlib.util.spec_from_file_location("schema", schema_file) + schema_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(schema_module) + + # Get the model class (assumes it follows naming convention) + model_name = f"{self.product_name.title().replace('_', '')}Experiment" + + if not hasattr(schema_module, model_name): + logger.warning(f"Schema class {model_name} not found") + return True + + model_class = getattr(schema_module, model_name) + + valid_count = 0 + error_count = 0 + + for _, row in data_df.iterrows(): + try: + # Convert row to dict + record = row.to_dict() + + # Validate against model + model_instance = model_class(**record) + valid_count += 1 + + except Exception as e: + error_count += 1 + if error_count <= 10: # Only show first 10 errors + logger.warning(f"Validation error for {row.get('name', 'unknown')}: {e}") + + logger.info(f"Validation: {valid_count} valid, {error_count} errors") + + return error_count == 0 + + def get_parquet_files(self) -> List[Path]: + """ + Get list of parquet files to upload + + Returns: + List of parquet file paths + """ + parquet_dir = self.product_dir / "parquet_data" + + if not parquet_dir.exists(): + return [] + + # Get all parquet files + parquet_files = list(parquet_dir.glob("*.parquet")) + + return parquet_files + + def setup_mpcontribs_project(self, dry_run: bool = False) -> bool: + """ + Setup empty MPContribs project with metadata only. + Does NOT upload data - that goes to S3. + + Args: + dry_run: If True, only simulate the operation + + Returns: + True if successful + """ + try: + # Convert config to dict if it's a Pydantic model + if hasattr(self.config, 'dict'): + config_dict = self.config.dict() + else: + config_dict = self.config + + return setup_mpcontribs_project( + product_config=config_dict, + dry_run=dry_run + ) + + except Exception as e: + logger.error(f"Error setting up MPContribs project: {e}") + import traceback + traceback.print_exc() + return False + + def upload_to_s3(self, dry_run: bool = True) -> bool: + """ + Upload parquet files directly to S3 OpenData. + + This is the new upload method that bypasses MPContribs library + and uploads parquet files directly to AWS S3. + + Args: + dry_run: If True, only simulate upload + + Returns: + True if successful + """ + parquet_dir = self.product_dir / "parquet_data" + + if not parquet_dir.exists(): + logger.error(f"Parquet directory not found: {parquet_dir}") + return False + + # Get parquet files (exclude very large ones by default) + parquet_files = list(parquet_dir.glob("*.parquet")) + + if not parquet_files: + logger.warning("No parquet files to upload") + return True + + # Show excluded fields info + if hasattr(self, 'schema_manager') and self.schema_manager: + logger.info("\n[Excluded from upload (embargoed):]") + for table in self.schema_manager.get_table_names(): + excluded = self.schema_manager.get_excluded_fields(table) + if excluded: + logger.info(f" {table}: {', '.join(excluded)}") + + # Show files that would be uploaded + large_files = {'temperature_logs.parquet', 'xrd_data_points.parquet'} + upload_files = [f for f in parquet_files if f.name not in large_files] + skipped_files = [f for f in parquet_files if f.name in large_files] + + total_size = sum(f.stat().st_size for f in upload_files) + + logger.info(f"\n{'[DRY RUN] ' if dry_run else ''}Files to upload ({len(upload_files)} files, {total_size / (1024*1024):.2f} MB):") + for pf in sorted(upload_files, key=lambda f: f.name): + size_kb = pf.stat().st_size / 1024 + logger.info(f" • {pf.name} ({size_kb:.1f} KB)") + + if skipped_files: + logger.info(f"\nSkipped (too large for OpenData):") + for pf in skipped_files: + size_mb = pf.stat().st_size / (1024 * 1024) + logger.info(f" • {pf.name} ({size_mb:.1f} MB)") + + if dry_run: + logger.info(f"\n[DRY RUN] Would upload {len(upload_files)} files to s3://materialsproject-contribs/alab_synthesis/{self.product_name}/") + return True + + # Perform actual upload + try: + uploaded = upload_product_to_s3( + product_name=self.product_name, + product_dir=self.product_dir, + metadata={'project': {'name': self.product_name, 'type': 'alab_synthesis'}}, + exclude_large=True, + dry_run=False + ) + + logger.info(f"✓ Uploaded {len(uploaded)} files to S3") + + # Record in state + if hasattr(self, 'state_manager') and self.state_manager: + self.state_manager.record_run( + run_type='s3_upload', + phases=['s3_upload'], + experiments=[self.product_name], + dry_run=False + ) + + return True + + except Exception as e: + logger.error(f"Failed to upload to S3: {e}") + return False + + def load_schemas(self): + """Load Pydantic schemas (schemas are now Python-first, not generated)""" + try: + from schema_manager import SchemaManager + self.schema_manager = SchemaManager() + + # List loaded schemas + schemas = self.schema_manager.get_table_names() + logger.info(f"✓ Loaded {len(schemas)} Pydantic schemas") + + # Show excluded fields + for table in schemas: + excluded = self.schema_manager.get_excluded_fields(table) + if excluded: + logger.info(f" {table}: {len(excluded)} fields excluded from upload ({', '.join(excluded)})") + + except Exception as e: + logger.warning(f"Could not load schemas: {e}") + self.schema_manager = None + + def run(self, + stages: List[str] = None, + dry_run: bool = True) -> bool: + """ + Run the complete pipeline + + Args: + stages: List of stages to run ['filter', 'transform', 'analyze', 'validate', 'diagram', 'mpcontribs', 'upload'] + dry_run: If True, simulate upload without actually uploading + + Returns: + True if successful + """ + if stages is None: + stages = ['filter', 'transform', 'analyze', 'validate', 'diagram', 'upload'] + + logger.info("=" * 60) + logger.info(f"Product Pipeline: {self.product_name}") + logger.info(f"Stages: {', '.join(stages)}") + logger.info(f"Mode: {'DRY RUN' if dry_run else 'LIVE'}") + if not dry_run and 'upload' in stages: + logger.info(f"Upload: MPContribs setup + S3 upload") + logger.info("=" * 60) + + # Load config + if not self.load_config(): + return False + + # Load Pydantic schemas (Python-first, no generation needed) + self.load_schemas() + + experiments_df = None + data_df = None + + # Stage 1: Filter experiments + if 'filter' in stages: + logger.info("\n[Stage 1/6] Filtering experiments...") + experiments_df = self.filter_experiments() + + if experiments_df.empty: + logger.error("No experiments match filter criteria") + return False + else: + # Load existing experiment list + exp_list_file = self.product_dir / "experiments.txt" + if exp_list_file.exists(): + exp_names = pd.read_csv(exp_list_file, header=None)[0].tolist() + experiments_df = pd.DataFrame({'name': exp_names}) + + # Stage 2: Transform to Parquet + if 'transform' in stages and experiments_df is not None: + logger.info("\n[Stage 2/6] Transforming to Parquet...") + if not self.transform_to_parquet(experiments_df): + return False + + # Validate schema after transformation + logger.info("\n[Stage 2.5/6] Validating schema...") + self._validate_schema() + + # Load data from consolidated experiments.parquet + data_df = self.load_experiments_data() + + # Stage 3: Run analyses + if 'analyze' in stages and experiments_df is not None: + logger.info("\n[Stage 3/6] Running analyses...") + analysis_results = self.run_analyses(experiments_df) + + # Merge analysis results with experiment data if any + if len(analysis_results) > 0: + # Merge on experiment name + data_df = data_df.merge( + analysis_results, + left_on='name', + right_on='experiment_name', + how='left' + ) + + # Stage 4: Validate + if 'validate' in stages and data_df is not None and len(data_df) > 0: + logger.info("\n[Stage 4/6] Validating with Pydantic...") + if not self.validate_data(data_df): + logger.warning("Validation errors found (continuing anyway)") + + # Stage 5: Generate Diagram + if 'diagram' in stages: + logger.info("\n[Stage 5/7] Generating schema diagram...") + self.generate_diagram() + + # Stage 6: Setup MPContribs (always when uploading for real) + if 'upload' in stages and not dry_run: + logger.info("\n[Stage 6/7] Setting up MPContribs project...") + if not self.setup_mpcontribs_project(dry_run=dry_run): + logger.warning("MPContribs setup failed (continuing anyway)") + + # Stage 7: Upload to S3 + if 'upload' in stages: + if dry_run: + logger.info("\n[Stage 7/7] Uploading to S3 OpenData (DRY RUN - no actual upload)...") + else: + logger.info("\n[Stage 7/7] Uploading to S3 OpenData...") + + if not self.upload_to_s3(dry_run=dry_run): + return False + + logger.info("\n" + "=" * 60) + logger.info("✓ Pipeline complete!") + logger.info("=" * 60) + + return True + + def _validate_schema(self): + """ + Validate parquet data using Pydantic schemas. + + Uses SchemaManager to load Pydantic schemas (Python-first, no YAML). + """ + from schema_validator import SchemaValidator + + parquet_dir = self.product_dir / "parquet_data" + + # Use pre-loaded schema manager or create new one + if not hasattr(self, 'schema_manager') or self.schema_manager is None: + from schema_manager import SchemaManager + self.schema_manager = SchemaManager() + + table_names = self.schema_manager.get_table_names() + + logger.info(f"✓ Validating {len(table_names)} tables using Pydantic schemas") + logger.info(f" Tables: {', '.join(table_names)}") + + try: + validator = SchemaValidator(self.schema_manager) + is_valid, errors, warnings = validator.validate_parquet_data(parquet_dir) + + if warnings: + from rich.console import Console + console = Console() + console.print("\n[yellow]⚠ Schema Validation Warnings:[/yellow]") + for warning in warnings[:20]: # Show first 20 + console.print(f" • {warning}") + if len(warnings) > 20: + console.print(f" • ... and {len(warnings) - 20} more") + + if errors: + from rich.console import Console + console = Console() + console.print("\n[red]✗ Schema Validation Errors:[/red]") + for error in errors: + console.print(f" • {error}") + + if is_valid: + if warnings: + from rich.console import Console + console = Console() + console.print("\n[green]✓ Schema validation passed with warnings[/green]") + else: + logger.info("✓ Schema validation passed") + else: + logger.error("✗ Schema validation failed") + + except Exception as e: + logger.warning(f"Schema validation skipped: {e}") + + def generate_diagram(self) -> bool: + """ + Generate schema diagram for this product's parquet data. + + Creates a markdown file with: + - Table overview with row counts + - Column details for each table + - Relationship diagram (Mermaid ERD) + + Output: {product_dir}/SCHEMA_DIAGRAM.md + """ + parquet_dir = self.product_dir / "parquet_data" + output_file = self.product_dir / "SCHEMA_DIAGRAM.md" + + if not parquet_dir.exists(): + logger.warning("No parquet data found, skipping diagram generation") + return False + + try: + from generate_diagram import ParquetSchemaAnalyzer, DiagramGenerator + + logger.info(f"Generating schema diagram for {self.product_name}...") + + # Analyze schema + analyzer = ParquetSchemaAnalyzer(parquet_dir) + analysis = analyzer.analyze() + + # Generate diagram + generator = DiagramGenerator(analysis) + + # Build output content + output_lines = [] + + # Header + output_lines.append(f"# {self.product_name} Schema Diagram") + output_lines.append("") + output_lines.append(f"**Auto-generated by A-Lab Pipeline**") + output_lines.append(f"") + output_lines.append(f"- **Product**: {self.product_name}") + output_lines.append(f"- **Tables**: {analysis['table_count']}") + output_lines.append(f"- **Total Rows**: {analysis['total_rows']:,}") + output_lines.append(f"- **Generated**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") + output_lines.append("") + output_lines.append("---") + output_lines.append("") + + # Add summary + output_lines.append(generator.generate_summary()) + + # Add Mermaid ERD + output_lines.append("\n---\n") + output_lines.append("## Entity Relationship Diagram\n") + output_lines.append(generator.generate_mermaid()) + + # Write to file + with open(output_file, 'w') as f: + f.write("\n".join(output_lines)) + + logger.info(f"✓ Generated schema diagram: {output_file}") + return True + + except ImportError: + logger.warning("Diagram generator not available (missing generate_diagram.py)") + return False + except Exception as e: + logger.error(f"Failed to generate diagram: {e}") + return False + + +def main(): + """CLI for running product pipeline""" + parser = argparse.ArgumentParser(description='Run product pipeline') + parser.add_argument('product_name', help='Name of the product') + parser.add_argument('--stages', nargs='+', + choices=['filter', 'transform', 'analyze', 'validate', 'diagram', 'upload'], + default=['filter', 'transform', 'analyze', 'validate', 'diagram', 'upload'], + help='Stages to run') + parser.add_argument('--live', action='store_true', + help='Run in live mode (actually upload to MPContribs)') + parser.add_argument('--dry-run', action='store_true', default=True, + help='Run in dry-run mode (default)') + + args = parser.parse_args() + + dry_run = not args.live + + pipeline = ProductPipeline(args.product_name) + success = pipeline.run(stages=args.stages, dry_run=dry_run) + + sys.exit(0 if success else 1) + + +if __name__ == '__main__': + main() diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/run_pipeline.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/run_pipeline.py new file mode 100755 index 000000000..18bcb5aef --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/run_pipeline.py @@ -0,0 +1,429 @@ +#!/usr/bin/env python3 +""" +A-Lab Pipeline Runner - Product-Based Pipeline + +Main CLI for managing data products and running pipelines. + +Commands: + create - Create a new data product interactively + list - List available products + run - Run pipeline for a product + status - Show product pipeline status + validate - Validate product configuration + +Usage: + python run_pipeline.py create + python run_pipeline.py run --product reaction_genome + python run_pipeline.py run --product reaction_genome --upload +""" + +import argparse +import sys +import subprocess +from pathlib import Path +from datetime import datetime +import logging +import warnings +from rich.console import Console +from rich.table import Table +from rich.panel import Panel +import pandas as pd + +# Suppress Pydantic config warnings +warnings.filterwarnings('ignore', message='Valid config keys have changed in V2') +warnings.filterwarnings('ignore', category=UserWarning, module='pydantic') + +# Add parent directories to path +sys.path.insert(0, str(Path(__file__).parent.parent)) +sys.path.insert(0, str(Path(__file__).parent.parent / "products")) + +from base_product import ProductManager, ProductConfig +from product_pipeline import ProductPipeline +from pipeline_state import PipelineStateManager + +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(levelname)s - %(message)s' +) +logger = logging.getLogger(__name__) + +console = Console() + + +class PipelineCLI: + """Main CLI for pipeline operations""" + + def __init__(self): + self.product_manager = ProductManager(Path("data/products")) + self.state_manager = PipelineStateManager() + + def cmd_create(self, args): + """Create a new data product interactively""" + console.print("\n[bold cyan]Creating New Data Product[/bold cyan]\n") + + try: + config = self.product_manager.create_product_interactive() + + console.print(f"\n[green]✓ Product '{config.name}' created successfully![/green]") + console.print(f"\nConfiguration saved to: data/products/{config.name}/config.yaml") + console.print(f"Pydantic schema: data/products/{config.name}/schema.py") + + # Automatically run dry run + console.print(f"\n[bold cyan]Running pipeline dry run...[/bold cyan]\n") + + from product_pipeline import ProductPipeline + pipeline = ProductPipeline(config.name) + success = pipeline.run(dry_run=True) + + if not success: + console.print(f"\n[yellow]⚠ Dry run encountered errors[/yellow]") + console.print(f"\n[dim]Fix issues and run manually:[/dim]") + console.print(f" ./run_product_pipeline.sh run --product {config.name}") + return 1 + + console.print(f"\n[green]✓ Dry run completed successfully![/green]") + + # Ask if they want to upload for real (MPContribs + S3) + import inquirer + upload = inquirer.confirm( + message=f"Upload {config.name}? (MPContribs setup + S3 upload)", + default=False + ) + + if upload: + console.print(f"\n[bold cyan]Uploading to MPContribs...[/bold cyan]\n") + pipeline = ProductPipeline(config.name) + success = pipeline.run(dry_run=False) + + if success: + console.print(f"\n[green]✓ Successfully uploaded to MPContribs![/green]") + else: + console.print(f"\n[red]✗ Upload failed[/red]") + return 1 + else: + console.print(f"\n[dim]To upload later, run:[/dim]") + console.print(f" ./run_product_pipeline.sh run --product {config.name} --upload") + + except KeyboardInterrupt: + console.print("\n[yellow]Product creation cancelled[/yellow]") + except Exception as e: + console.print(f"\n[red]Error creating product: {e}[/red]") + import traceback + console.print(traceback.format_exc()) + return 1 + + return 0 + + def cmd_list(self, args): + """List available data products""" + products = self.product_manager.list_products() + + if not products: + console.print("[yellow]No data products found[/yellow]") + console.print("\nCreate one with: python run_pipeline.py create") + return 0 + + # Create table + table = Table(title="Available Data Products") + table.add_column("Product", style="cyan") + table.add_column("Experiments", justify="right") + table.add_column("Analyses", justify="right") + table.add_column("Last Run", style="dim") + table.add_column("Status") + + for product_name in products: + try: + # Load config + config = self.product_manager.get_product_config(product_name) + + # Get experiment count + exp_file = Path(f"data/products/{product_name}/experiments.txt") + exp_count = 0 + if exp_file.exists(): + exp_count = len(exp_file.read_text().strip().split('\n')) + + # Get enabled analyses + enabled_analyses = sum(1 for a in config.analyses if a.get('enabled', True)) + + # Get last run from state + last_run = "Never" + status = "Not run" + + # Check for recent runs in state + if len(self.state_manager.runs) > 0: + product_runs = self.state_manager.runs[ + self.state_manager.runs['experiment_name'].str.contains(product_name, na=False) + ] + if len(product_runs) > 0: + last_timestamp = product_runs['run_timestamp'].max() + last_run = pd.to_datetime(last_timestamp).strftime('%Y-%m-%d %H:%M') + + # Get status + last_status = product_runs[ + product_runs['run_timestamp'] == last_timestamp + ]['status'].iloc[0] + + if last_status == 'success': + status = "[green]✓ Success[/green]" + elif last_status == 'failed': + status = "[red]✗ Failed[/red]" + else: + status = "[yellow]⟳ Running[/yellow]" + + table.add_row( + product_name, + str(exp_count), + str(enabled_analyses), + last_run, + status + ) + + except Exception as e: + table.add_row(product_name, "?", "?", "?", f"[red]Error[/red]") + + console.print(table) + + # Show additional info + console.print(f"\n[dim]Products directory: data/products/[/dim]") + console.print(f"[dim]Run a product: python run_pipeline.py run --product [/dim]") + + return 0 + + def cmd_run(self, args): + """Run pipeline for a data product""" + if not args.product: + console.print("[red]Error: --product is required[/red]") + return 1 + + # Check if product exists + products = self.product_manager.list_products() + if args.product not in products: + console.print(f"[red]Product '{args.product}' not found[/red]") + console.print(f"\nAvailable products: {', '.join(products)}") + return 1 + + # Run pipeline + pipeline = ProductPipeline(args.product) + + dry_run = not args.upload + + success = pipeline.run( + stages=args.stages, + dry_run=dry_run + ) + + return 0 if success else 1 + + def cmd_status(self, args): + """Show status for a product or all products""" + if args.product: + products = [args.product] + else: + products = self.product_manager.list_products() + + if not products: + console.print("[yellow]No products found[/yellow]") + return 0 + + for product_name in products: + console.print(Panel(f"[bold]{product_name}[/bold]", expand=False)) + + try: + # Load config + config = self.product_manager.get_product_config(product_name) + product_dir = Path(f"data/products/{product_name}") + + # Check what exists + checks = { + 'Configuration': (product_dir / 'config.yaml').exists(), + 'Pydantic Schema': (product_dir / 'schema.py').exists(), + 'Experiment List': (product_dir / 'experiments.txt').exists(), + 'Parquet Data': (product_dir / 'parquet_data').exists(), + 'Analysis Results': (product_dir / 'analysis_results').exists(), + } + + for item, exists in checks.items(): + status = "[green]✓[/green]" if exists else "[red]✗[/red]" + console.print(f" {status} {item}") + + # Show experiment count + exp_file = product_dir / 'experiments.txt' + if exp_file.exists(): + exp_count = len(exp_file.read_text().strip().split('\n')) + console.print(f"\n Experiments: {exp_count}") + + # Show filter info + if config.experiment_filter.types: + console.print(f" Types: {', '.join(config.experiment_filter.types)}") + + # Show analyses + enabled = [a['name'] for a in config.analyses if a.get('enabled', True)] + if enabled: + console.print(f" Analyses: {', '.join(enabled)}") + + console.print() + + except Exception as e: + console.print(f" [red]Error: {e}[/red]\n") + + return 0 + + def cmd_validate(self, args): + """Validate product configuration""" + if not args.product: + console.print("[red]Error: --product is required[/red]") + return 1 + + try: + # Load config + config = self.product_manager.get_product_config(args.product) + + console.print(f"\n[bold]Validating {args.product}[/bold]\n") + + # Validate configuration structure + issues = [] + warnings = [] + + # Check required fields + if not config.experiment_filter.types and not config.experiment_filter.experiment_names: + issues.append("No experiment filter defined (types or experiment_names required)") + + if not config.data_schema: + warnings.append("No schema fields defined") + + # Check analyses exist + from analyses.base_analyzer import AnalysisPluginManager + plugin_manager = AnalysisPluginManager() + available = plugin_manager.list_analyzers() + + for analysis in config.analyses: + if analysis['name'] not in available: + warnings.append(f"Analysis '{analysis['name']}' not found in available analyzers") + + # Check Pydantic schema + schema_file = Path(f"data/products/{args.product}/schema.py") + if schema_file.exists(): + console.print(" [green]✓[/green] Pydantic schema exists") + + # Try to import it + try: + import importlib.util + spec = importlib.util.spec_from_file_location("schema", schema_file) + schema_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(schema_module) + console.print(" [green]✓[/green] Pydantic schema is valid Python") + except Exception as e: + issues.append(f"Pydantic schema error: {e}") + else: + warnings.append("No Pydantic schema generated") + + # Report results + if issues: + console.print("\n[red]Issues found:[/red]") + for issue in issues: + console.print(f" • {issue}") + + if warnings: + console.print("\n[yellow]Warnings:[/yellow]") + for warning in warnings: + console.print(f" • {warning}") + + if not issues and not warnings: + console.print("[green]✓ Configuration is valid![/green]") + + return 1 if issues else 0 + + except Exception as e: + console.print(f"[red]Error validating product: {e}[/red]") + return 1 + + def cmd_regenerate(self, args): + """Show schema info (schemas are now Python-first, no regeneration needed)""" + console.print(f"\n[bold]Pydantic Schemas (Python-first)[/bold]\n") + console.print("[cyan]Schemas are now defined directly in Python at data/products/schema/[/cyan]") + console.print("[cyan]No regeneration needed - just edit the .py files directly.[/cyan]\n") + + try: + from schema_manager import SchemaManager + schema_manager = SchemaManager() + + console.print("[bold]Available Schemas:[/bold]") + for table_name in schema_manager.get_table_names(): + schema = schema_manager.get_schema(table_name) + excluded = schema_manager.get_excluded_fields(table_name) + + console.print(f"\n [green]{table_name}[/green] ({schema.__name__})") + console.print(f" Fields: {len(schema.model_fields)}") + if excluded: + console.print(f" [yellow]Excluded from upload: {', '.join(excluded)}[/yellow]") + + return 0 + + except Exception as e: + console.print(f"[red]Error loading schemas: {e}[/red]") + return 1 + + +def main(): + parser = argparse.ArgumentParser( + description='A-Lab Pipeline Runner - Product-Based Pipeline' + ) + + subparsers = parser.add_subparsers(dest='command', help='Commands') + + # Create command + create_parser = subparsers.add_parser('create', help='Create new data product') + + # List command + list_parser = subparsers.add_parser('list', help='List available products') + + # Run command + run_parser = subparsers.add_parser('run', help='Run pipeline for a product') + run_parser.add_argument('--product', '-p', required=True, help='Product name') + run_parser.add_argument('--stages', '-s', nargs='+', + choices=['filter', 'transform', 'analyze', 'validate', 'diagram', 'upload'], + help='Stages to run (default: all)') + run_parser.add_argument('--upload', action='store_true', + help='Upload for real: MPContribs setup + S3 upload (default is dry run)') + + # Status command + status_parser = subparsers.add_parser('status', help='Show product status') + status_parser.add_argument('--product', '-p', help='Product name (or all)') + + # Validate command + val_parser = subparsers.add_parser('validate', help='Validate product config') + val_parser.add_argument('--product', '-p', required=True, help='Product name') + + # Regenerate command + regen_parser = subparsers.add_parser('regenerate', help='Regenerate Pydantic schema from config') + regen_parser.add_argument('--product', '-p', required=True, help='Product name') + + args = parser.parse_args() + + if not args.command: + # Default to 'create' if no command specified + console.print("[cyan]No command specified, starting product creation...[/cyan]\n") + args.command = 'create' + + # Execute command + cli = PipelineCLI() + + if args.command == 'create': + return cli.cmd_create(args) + elif args.command == 'list': + return cli.cmd_list(args) + elif args.command == 'run': + return cli.cmd_run(args) + elif args.command == 'status': + return cli.cmd_status(args) + elif args.command == 'validate': + return cli.cmd_validate(args) + elif args.command == 'regenerate': + return cli.cmd_regenerate(args) + else: + parser.print_help() + return 1 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/s3_uploader.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/s3_uploader.py new file mode 100644 index 000000000..6fa2c042b --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/s3_uploader.py @@ -0,0 +1,420 @@ +#!/usr/bin/env python3 +""" +S3 Uploader for MPContribs OpenData + +Direct upload of parquet files to AWS S3 for MPContribs. +Based on MPContribs documentation for large data uploads. + +Requires AWS credentials from MP staff: +- aws_access_key_id +- aws_secret_access_key + +Set in .env file: + aws_access_key_id=your_key_here + aws_secret_access_key=your_secret_here +""" + +import json +import logging +import os +import sys +from pathlib import Path +from typing import Optional, List, Dict, Any + +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq + +# Import config loader +sys.path.insert(0, str(Path(__file__).parent.parent / "config")) +from config_loader import get_config + +logger = logging.getLogger(__name__) + +# Load configuration (env vars > yaml > defaults) +config = get_config() + + +def prepare_metadata_from_config(config: Dict[str, Any]) -> Dict[str, Dict[str, Any]]: + """ + Prepare metadata dict from product configuration for embedding in parquet. + + Args: + config: Product configuration dictionary (from config.yaml or ProductConfig) + + Returns: + Dict formatted for parquet schema metadata embedding + + Example: + >>> config = {'name': 'my_product', 'metadata': {'title': 'My Dataset'}} + >>> metadata = prepare_metadata_from_config(config) + >>> # Use with S3Uploader.upload_parquet(metadata=metadata) + """ + metadata = {} + + # Project info + project_info = { + 'name': config.get('name', 'alab_synthesis'), + 'type': 'alab_synthesis' + } + + # Add metadata if available + if 'metadata' in config: + meta = config['metadata'] + if 'title' in meta: + project_info['title'] = meta['title'] + if 'description' in meta: + project_info['description'] = meta['description'] + if 'authors' in meta: + project_info['authors'] = meta['authors'] + + metadata['project'] = project_info + + # Add column metadata if schema is available + if 'data_schema' in config: + columns = {} + for field_name, field_def in config['data_schema'].items(): + col_meta = {} + + # Add description + if isinstance(field_def, dict): + if 'description' in field_def: + col_meta['description'] = field_def['description'] + if 'unit' in field_def: + col_meta['unit'] = field_def['unit'] + elif hasattr(field_def, 'description'): + col_meta['description'] = field_def.description + if hasattr(field_def, 'unit'): + col_meta['unit'] = field_def.unit + + if col_meta: + columns[field_name] = col_meta + + if columns: + metadata['columns'] = columns + + return metadata + + +class S3Uploader: + """Upload parquet files to AWS S3 OpenData for MPContribs.""" + + def __init__(self, project_name: str, api_key_id: str = None, api_key_secret: str = None): + """ + Initialize S3 uploader. + + Args: + project_name: MPContribs project name (e.g., 'alab_synthesis') + api_key_id: AWS access key ID (or from env AWS_ACCESS_KEY_ID) + api_key_secret: AWS secret access key (or from env AWS_SECRET_ACCESS_KEY) + """ + self.project_name = project_name + + # Load credentials from env or args + try: + from dotenv import load_dotenv + env_file = Path(__file__).parent.parent.parent / ".env" + if env_file.exists(): + load_dotenv(env_file) + except ImportError: + pass + + self.api_key_id = api_key_id or os.environ.get('aws_access_key_id') or os.environ.get('AWS_ACCESS_KEY_ID') + self.api_key_secret = api_key_secret or os.environ.get('aws_secret_access_key') or os.environ.get('AWS_SECRET_ACCESS_KEY') + + self._client = None + + @property + def client(self): + """Lazy load S3 client.""" + if self._client is None: + if not self.api_key_id or not self.api_key_secret: + raise ValueError( + "AWS credentials not set. Add to .env file:\n" + " aws_access_key_id=your_key\n" + " aws_secret_access_key=your_secret" + ) + + try: + import boto3 + self._client = boto3.client( + "s3", + aws_access_key_id=self.api_key_id, + aws_secret_access_key=self.api_key_secret, + ) + logger.info("Connected to AWS S3") + except ImportError: + raise ImportError("boto3 not installed. Run: pip install boto3") + + return self._client + + def upload_parquet( + self, + local_path: Path, + key: str = None, + metadata: Dict = None, + dry_run: bool = True + ) -> str: + """ + Upload a parquet file to S3. + + Args: + local_path: Path to local parquet file + key: S3 key (default: {project_name}/{filename}) + metadata: Optional metadata to embed in parquet schema + dry_run: If True, don't actually upload + + Returns: + S3 URL of uploaded file + """ + local_path = Path(local_path) + key = key or f"{self.project_name}/{local_path.name}" + + s3_url = f"s3://{config.s3_bucket}/{key}" + + if dry_run: + size_mb = local_path.stat().st_size / (1024 * 1024) + logger.info(f"[DRY RUN] Would upload: {local_path.name} ({size_mb:.2f} MB) -> {s3_url}") + return s3_url + + # If metadata provided, embed it in the parquet file + if metadata: + table = pq.read_table(local_path) + new_metadata = { + **table.schema.metadata, + **{ + f"{key}.{sub_key}": json.dumps(v) + for key, vals in metadata.items() + for sub_key, v in vals.items() + }, + } + table = table.replace_schema_metadata(metadata=new_metadata) + + # Write to temp file with metadata + temp_path = local_path.with_suffix('.tmp.parquet') + pq.write_table(table, temp_path) + upload_path = temp_path + else: + upload_path = local_path + + try: + with open(upload_path, "rb") as f: + self.client.upload_fileobj(f, Bucket=config.s3_bucket, Key=key) + + logger.info(f"✓ Uploaded: {local_path.name} -> {s3_url}") + return s3_url + + finally: + # Clean up temp file + if metadata and temp_path.exists(): + temp_path.unlink() + + def upload_all_parquet( + self, + parquet_dir: Path, + metadata: Dict = None, + exclude_large: bool = True, + dry_run: bool = True + ) -> Dict[str, str]: + """ + Upload all parquet files from a directory. + + Args: + parquet_dir: Directory containing parquet files + metadata: Optional metadata to embed in all files + exclude_large: If True, skip temperature_logs and xrd_data_points (very large) + dry_run: If True, don't actually upload + + Returns: + Dict of filename -> S3 URL + """ + parquet_dir = Path(parquet_dir) + + if not parquet_dir.exists(): + logger.error(f"Directory not found: {parquet_dir}") + return {} + + # Find all parquet files + parquet_files = list(parquet_dir.glob("*.parquet")) + + if not parquet_files: + logger.warning(f"No parquet files found in {parquet_dir}") + return {} + + # Optionally exclude large files + large_files = {'temperature_logs.parquet', 'xrd_data_points.parquet'} + if exclude_large: + parquet_files = [f for f in parquet_files if f.name not in large_files] + logger.info(f"Excluding large files: {large_files}") + + # Sort by size (smallest first) + parquet_files.sort(key=lambda f: f.stat().st_size) + + # Calculate total size + total_size = sum(f.stat().st_size for f in parquet_files) + logger.info(f"{'[DRY RUN] ' if dry_run else ''}Uploading {len(parquet_files)} files ({total_size / (1024*1024):.2f} MB total)") + + # Upload each file + uploaded = {} + for pf in parquet_files: + s3_url = self.upload_parquet(pf, metadata=metadata, dry_run=dry_run) + uploaded[pf.name] = s3_url + + return uploaded + + def delete_file(self, key: str, dry_run: bool = True) -> bool: + """ + Delete a file from S3. + + Args: + key: S3 key (e.g., 'project_name/file.parquet') + dry_run: If True, don't actually delete + + Returns: + True if successful + """ + if dry_run: + logger.info(f"[DRY RUN] Would delete: s3://{config.s3_bucket}/{key}") + return True + + try: + self.client.delete_object(Bucket=config.s3_bucket, Key=key) + logger.info(f"✓ Deleted: s3://{config.s3_bucket}/{key}") + return True + except Exception as e: + logger.error(f"Failed to delete {key}: {e}") + return False + + def list_files(self, details: bool = False) -> List[str] | List[Dict]: + """ + List all files in the project's S3 prefix. + + Args: + details: If True, return detailed info (size, last modified, etc.) + + Returns: + List of file keys (strings) or list of dicts with details + """ + try: + response = self.client.list_objects_v2( + Bucket=config.s3_bucket, + Prefix=f"{self.project_name}/" + ) + + if details: + files = [] + for obj in response.get('Contents', []): + files.append({ + 'key': obj['Key'], + 'size_bytes': obj['Size'], + 'size_mb': obj['Size'] / (1024 * 1024), + 'last_modified': obj['LastModified'], + 'etag': obj['ETag'] + }) + return files + else: + files = [] + for obj in response.get('Contents', []): + files.append(obj['Key']) + return files + + except Exception as e: + logger.error(f"Failed to list files: {e}") + return [] + + def get_upload_url(self, key: str) -> str: + """ + Get public HTTPS URL for accessing uploaded file. + + Args: + key: S3 key (e.g., 'alab_synthesis/product/file.parquet') + + Returns: + HTTPS URL for accessing the file + """ + return f"https://{config.s3_bucket}.s3.amazonaws.com/{key}" + + +def upload_product_to_s3( + product_name: str, + product_dir: Path, + metadata: Dict = None, + exclude_large: bool = True, + dry_run: bool = True, + auto_metadata: bool = True +) -> Dict[str, str]: + """ + Upload a product's parquet files to S3. + + Args: + product_name: Name of the product (used as S3 prefix) + product_dir: Path to product directory + metadata: Optional metadata to embed (if None and auto_metadata=True, reads from config.yaml) + exclude_large: If True, skip very large files (temperature_logs, xrd_data_points) + dry_run: If True, don't actually upload + auto_metadata: If True and metadata is None, try to load from config.yaml + + Returns: + Dict of filename -> S3 URL + + Example: + >>> from pathlib import Path + >>> uploaded = upload_product_to_s3( + ... product_name='reaction_genome', + ... product_dir=Path('data/products/reaction_genome'), + ... dry_run=False + ... ) + """ + uploader = S3Uploader(project_name=f"alab_synthesis/{product_name}") + parquet_dir = product_dir / "parquet_data" + + # Auto-load metadata from config.yaml if requested + if metadata is None and auto_metadata: + config_file = product_dir / "config.yaml" + if config_file.exists(): + try: + import yaml + with open(config_file, 'r') as f: + config = yaml.safe_load(f) + metadata = prepare_metadata_from_config(config) + logger.info(f"Loaded metadata from {config_file.name}") + except Exception as e: + logger.warning(f"Could not load metadata from config: {e}") + + return uploader.upload_all_parquet( + parquet_dir=parquet_dir, + metadata=metadata, + exclude_large=exclude_large, + dry_run=dry_run + ) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Upload parquet files to MPContribs S3") + parser.add_argument("--project", "-p", required=True, help="MPContribs project name") + parser.add_argument("--dir", "-d", required=True, help="Directory containing parquet files") + parser.add_argument("--upload", action="store_true", help="Actually upload (default is dry run)") + parser.add_argument("--include-large", action="store_true", help="Include temperature_logs and xrd_data_points") + + args = parser.parse_args() + + dry_run = not args.upload + + if dry_run: + print("=" * 60) + print("DRY RUN MODE - No files will be uploaded") + print("Use --upload to actually upload to S3") + print("=" * 60) + print() + + uploader = S3Uploader(project_name=args.project) + uploaded = uploader.upload_all_parquet( + parquet_dir=Path(args.dir), + exclude_large=not args.include_large, + dry_run=dry_run + ) + + print(f"\n{'Would upload' if dry_run else 'Uploaded'} {len(uploaded)} files") + diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/test_integrated_pipeline.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/test_integrated_pipeline.py new file mode 100755 index 000000000..bffc8cc0a --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/test_integrated_pipeline.py @@ -0,0 +1,920 @@ +#!/usr/bin/env python3 +""" +Comprehensive Integration Tests for A-Lab Pipeline + +Tests the complete pipeline including: +- Auto-discovery (schemas, analyses) +- Filter configurations +- Hook system +- Edge cases and error handling +- End-to-end workflow + +Usage: + python test_integrated_pipeline.py # Run all tests + python test_integrated_pipeline.py --verbose # Verbose output + pytest test_integrated_pipeline.py # Use pytest +""" + +import sys +import os +import tempfile +import shutil +from pathlib import Path +from datetime import datetime +import json +import traceback + +# Add paths +sys.path.insert(0, str(Path("data/products"))) +sys.path.insert(0, str(Path("data/pipeline"))) +sys.path.insert(0, str(Path("data/analyses"))) +sys.path.insert(0, str(Path("data"))) +sys.path.insert(0, str(Path("data/tools"))) + +# Test results tracking +test_results = [] + + +def log_test(test_name: str, passed: bool, message: str = ""): + """Track test results""" + test_results.append({ + 'test': test_name, + 'passed': passed, + 'message': message + }) + status = "✓" if passed else "✗" + print(f" {status} {test_name}") + if message and not passed: + print(f" {message}") + + +def test_header(title: str): + """Print test section header""" + print("\n" + "=" * 70) + print(f" {title}") + print("=" * 70) + + +# ============================================================================= +# Auto-Discovery Tests +# ============================================================================= + +def test_schema_auto_discovery(): + """Test that schemas are auto-discovered from schema directory""" + test_header("TEST 1: Schema Auto-Discovery") + + try: + from schema_manager import SchemaManager + + manager = SchemaManager() + discovered = manager.get_table_names() + + # Should discover at least these core schemas + expected_schemas = ['experiments', 'experiment_elements', 'powder_doses'] + + all_found = all(s in discovered for s in expected_schemas) + log_test( + "Core schemas discovered", + all_found, + f"Expected {expected_schemas}, found {discovered}" + ) + + # Test that schemas have model_fields + for schema_name in discovered: + schema_class = manager.get_schema(schema_name) + has_fields = hasattr(schema_class, 'model_fields') and len(schema_class.model_fields) > 0 + log_test( + f"Schema '{schema_name}' has fields", + has_fields, + f"Fields: {len(schema_class.model_fields) if has_fields else 0}" + ) + + # Test excluded fields detection + experiments_schema = manager.get_schema('experiments') + if experiments_schema: + excluded = manager.get_excluded_fields('experiments') + log_test( + "Excluded fields detected", + True, + f"Found {len(excluded)} excluded fields: {excluded}" + ) + + return True + + except Exception as e: + log_test("Schema auto-discovery", False, str(e)) + traceback.print_exc() + return False + + +def test_analysis_auto_discovery(): + """Test that analyses are auto-discovered from analyses directory""" + test_header("TEST 2: Analysis Auto-Discovery") + + try: + from base_analyzer import AnalysisPluginManager + + manager = AnalysisPluginManager() + discovered = manager.list_analyzers() + + # Should discover at least these built-in analyses + expected_analyses = ['xrd_dara', 'powder_statistics'] + + all_found = all(a in discovered for a in expected_analyses) + log_test( + "Built-in analyses discovered", + all_found, + f"Expected {expected_analyses}, found {discovered}" + ) + + # Test analyzer metadata + for analyzer_name in discovered: + analyzer = manager.get_analyzer(analyzer_name) + + has_analyze = hasattr(analyzer, 'analyze') and callable(analyzer.analyze) + log_test(f"Analyzer '{analyzer_name}' has analyze()", has_analyze) + + has_schema = hasattr(analyzer, 'get_output_schema') and callable(analyzer.get_output_schema) + log_test(f"Analyzer '{analyzer_name}' has get_output_schema()", has_schema) + + if has_schema: + schema = analyzer.get_output_schema() + log_test( + f"Analyzer '{analyzer_name}' schema valid", + isinstance(schema, dict) and len(schema) > 0, + f"Fields: {list(schema.keys())}" + ) + + return True + + except Exception as e: + log_test("Analysis auto-discovery", False, str(e)) + traceback.print_exc() + return False + + +def test_custom_schema_hook(): + """Test adding a custom schema via hook""" + test_header("TEST 3: Custom Schema Hook") + + try: + from schema_manager import SchemaManager + from pydantic import BaseModel, Field + + # Create a temporary custom schema + custom_schema_dir = Path("data/products/schema") + custom_schema_file = custom_schema_dir / "test_custom_schema.py" + + # Write custom schema + custom_schema_code = '''"""Test custom schema""" +from pydantic import BaseModel, Field + +class TestCustomData(BaseModel, extra="forbid"): + """Test custom data schema""" + __schema_table__ = "test_custom_data" + + experiment_id: str = Field(description="Experiment ID") + custom_field: float = Field(description="Custom measurement") +''' + + with open(custom_schema_file, 'w') as f: + f.write(custom_schema_code) + + try: + # Re-discover schemas + manager = SchemaManager() + discovered = manager.get_table_names() + + custom_found = 'test_custom_data' in discovered + log_test( + "Custom schema auto-discovered", + custom_found, + f"Found in: {discovered}" + ) + + if custom_found: + custom_schema = manager.get_schema('test_custom_data') + log_test( + "Custom schema loaded correctly", + custom_schema is not None and hasattr(custom_schema, 'model_fields'), + f"Fields: {list(custom_schema.model_fields.keys()) if custom_schema else []}" + ) + + finally: + # Cleanup + if custom_schema_file.exists(): + custom_schema_file.unlink() + + return True + + except Exception as e: + log_test("Custom schema hook", False, str(e)) + traceback.print_exc() + return False + + +def test_custom_analyzer_hook(): + """Test adding a custom analyzer via hook""" + test_header("TEST 4: Custom Analyzer Hook") + + try: + from base_analyzer import AnalysisPluginManager + + # Create a temporary custom analyzer + custom_analyzer_file = Path("data/analyses/test_custom_analyzer.py") + + # Write custom analyzer + custom_analyzer_code = '''"""Test custom analyzer""" +from pathlib import Path +import pandas as pd +from base_analyzer import BaseAnalyzer + +class TestCustomAnalyzer(BaseAnalyzer): + """Test custom analysis""" + name = "test_custom" + description = "Test custom analysis for testing" + cli_flag = "--test-custom" + + def analyze(self, experiments_df: pd.DataFrame, parquet_dir: Path) -> pd.DataFrame: + results = [] + for _, exp in experiments_df.iterrows(): + results.append({ + 'experiment_name': exp.get('name', 'test'), + 'test_metric': 42.0 + }) + return pd.DataFrame(results) + + def get_output_schema(self): + return { + 'test_metric': {'type': 'float', 'required': True, 'description': 'Test metric'} + } +''' + + with open(custom_analyzer_file, 'w') as f: + f.write(custom_analyzer_code) + + try: + # Re-discover analyzers + manager = AnalysisPluginManager() + discovered = manager.list_analyzers() + + custom_found = 'test_custom' in discovered + log_test( + "Custom analyzer auto-discovered", + custom_found, + f"Found in: {discovered}" + ) + + if custom_found: + custom_analyzer = manager.get_analyzer('test_custom') + log_test( + "Custom analyzer loaded correctly", + custom_analyzer is not None and hasattr(custom_analyzer, 'analyze'), + f"Name: {custom_analyzer.name if custom_analyzer else 'N/A'}" + ) + + # Test analyzer can be instantiated + schema = custom_analyzer.get_output_schema() + log_test( + "Custom analyzer schema valid", + isinstance(schema, dict) and 'test_metric' in schema, + f"Schema: {schema}" + ) + + finally: + # Cleanup + if custom_analyzer_file.exists(): + custom_analyzer_file.unlink() + + return True + + except Exception as e: + log_test("Custom analyzer hook", False, str(e)) + traceback.print_exc() + return False + + +# ============================================================================= +# Filter Configuration Tests +# ============================================================================= + +def test_filter_configurations(): + """Test various filter configurations""" + test_header("TEST 5: Filter Configurations") + + try: + from base_product import ExperimentFilter + + # Test 1: Simple type filter + filter1 = ExperimentFilter(types=["NSC"]) + query1 = filter1.to_mongo_query() + log_test( + "Simple type filter", + 'name' in query1 and '$regex' in query1['name'], + f"Query: {query1}" + ) + + # Test 2: Multiple types + filter2 = ExperimentFilter(types=["NSC", "Na"]) + query2 = filter2.to_mongo_query() + log_test( + "Multiple type filter", + 'name' in query2, + f"Query: {query2}" + ) + + # Test 3: Status filter + filter3 = ExperimentFilter(status=["completed"]) + query3 = filter3.to_mongo_query() + log_test( + "Status filter", + 'status' in query3, + f"Query: {query3}" + ) + + # Test 4: XRD requirement + filter4 = ExperimentFilter(has_xrd=True) + query4 = filter4.to_mongo_query() + log_test( + "XRD requirement filter", + 'metadata.diffraction_results.sampleid_in_aeris' in query4, + f"Query: {query4}" + ) + + # Test 5: Combined filters + filter5 = ExperimentFilter( + types=["NSC"], + status=["completed"], + has_xrd=True + ) + query5 = filter5.to_mongo_query() + log_test( + "Combined filters", + len(query5) >= 2, + f"Query has {len(query5)} conditions" + ) + + # Test 6: Specific experiment names + filter6 = ExperimentFilter(experiment_names=["NSC_249_001", "NSC_249_002"]) + query6 = filter6.to_mongo_query() + log_test( + "Specific experiment names", + 'name' in query6 and '$in' in query6['name'], + f"Query: {query6}" + ) + + # Test 7: Date range filter + filter7 = ExperimentFilter( + date_range={ + "start": "2024-01-01", + "end": "2024-12-31" + } + ) + query7 = filter7.to_mongo_query() + log_test( + "Date range filter", + 'last_updated' in query7, + f"Query: {query7}" + ) + + # Test 8: Empty filter (should match all) + filter8 = ExperimentFilter() + query8 = filter8.to_mongo_query() + log_test( + "Empty filter (match all)", + query8 == {}, + f"Query: {query8}" + ) + + return True + + except Exception as e: + log_test("Filter configurations", False, str(e)) + traceback.print_exc() + return False + + +# ============================================================================= +# Edge Cases and Error Handling +# ============================================================================= + +def test_edge_cases(): + """Test edge cases and potential breaking scenarios""" + test_header("TEST 6: Edge Cases & Error Handling") + + try: + from base_product import ProductConfig, ExperimentFilter, ProductMetadata + + # Test 1: Invalid product name + try: + invalid_config = ProductConfig( + name="invalid name with spaces!", + experiment_filter=ExperimentFilter() + ) + log_test("Invalid product name rejected", False, "Should have raised ValueError") + except ValueError: + log_test("Invalid product name rejected", True, "Correctly rejected invalid name") + + # Test 2: Empty product name + try: + empty_config = ProductConfig( + name="", + experiment_filter=ExperimentFilter() + ) + log_test("Empty product name rejected", False, "Should have raised ValueError") + except (ValueError, Exception): + log_test("Empty product name rejected", True, "Correctly rejected empty name") + + # Test 3: Valid product with underscores + try: + valid_config = ProductConfig( + name="valid_product_123", + experiment_filter=ExperimentFilter() + ) + log_test("Valid product name accepted", True, f"Name: {valid_config.name}") + except Exception as e: + log_test("Valid product name accepted", False, str(e)) + + # Test 4: Missing required fields + try: + from schema_manager import SchemaManager + manager = SchemaManager() + exp_schema = manager.get_schema('experiments') + + if exp_schema: + # Try to create with missing required fields + try: + invalid_exp = exp_schema( + experiment_id="test", + name="test" + # Missing other required fields + ) + log_test("Missing required fields rejected", False, "Should have raised validation error") + except Exception: + log_test("Missing required fields rejected", True, "Correctly rejected") + except Exception as e: + log_test("Schema validation test", False, str(e)) + + # Test 5: Field validation (optional fields accept None) + try: + from schema_manager import SchemaManager + manager = SchemaManager() + exp_schema = manager.get_schema('experiments') + + if exp_schema: + # Test that validation works at all + try: + # Try to create with invalid status (not in Literal) + invalid_exp = exp_schema( + experiment_id="test", + name="test", + experiment_type="TEST", + target_formula="test", + last_updated=datetime.now(), + status="invalid_status", # Not in Literal["completed", "error", "active", "unknown"] + ) + log_test("Invalid enum values rejected", False, "Should have rejected invalid status") + except Exception: + log_test("Invalid enum values rejected", True, "Correctly rejected invalid status") + + # Test that valid data passes + try: + valid_exp = exp_schema( + experiment_id="test", + name="test", + experiment_type="TEST", + target_formula="test", + last_updated=datetime.now(), + status="completed", + heating_temperature=1100.0 # Valid temperature + ) + log_test("Valid data accepted", True, "Schema accepts valid data") + except Exception as e: + log_test("Valid data accepted", False, f"Schema rejected valid data: {e}") + except Exception as e: + log_test("Schema validation test", False, str(e)) + + # Test 6: Non-existent analyzer + try: + from base_analyzer import AnalysisPluginManager + manager = AnalysisPluginManager() + + non_existent = manager.get_analyzer("this_does_not_exist") + log_test( + "Non-existent analyzer handled", + non_existent is None, + "Should return None for missing analyzer" + ) + except Exception as e: + log_test("Non-existent analyzer test", False, str(e)) + + # Test 7: Non-existent schema + try: + from schema_manager import SchemaManager + manager = SchemaManager() + + non_existent = manager.get_schema("this_does_not_exist") + log_test( + "Non-existent schema handled", + non_existent is None, + "Should return None for missing schema" + ) + except Exception as e: + log_test("Non-existent schema test", False, str(e)) + + return True + + except Exception as e: + log_test("Edge cases", False, str(e)) + traceback.print_exc() + return False + + +def test_config_files(): + """Test configuration files exist and are valid""" + test_header("TEST 7: Configuration Files") + + try: + import yaml + + # Test defaults.yaml + defaults_file = Path("data/config/defaults.yaml") + if defaults_file.exists(): + with open(defaults_file) as f: + defaults = yaml.safe_load(f) + + log_test( + "defaults.yaml exists and valid", + isinstance(defaults, dict) and 'version' in defaults, + f"Version: {defaults.get('version')}" + ) + + # Check required sections + required_sections = ['mongodb', 'parquet', 'analyses', 'upload'] + for section in required_sections: + log_test( + f"defaults.yaml has '{section}' section", + section in defaults, + f"Keys: {list(defaults.keys())}" + ) + else: + log_test("defaults.yaml exists", False, "File not found") + + # Test filters.yaml + filters_file = Path("data/config/filters.yaml") + if filters_file.exists(): + with open(filters_file) as f: + filters = yaml.safe_load(f) + + log_test( + "filters.yaml exists and valid", + isinstance(filters, dict), + f"Keys: {list(filters.keys())}" + ) + + # Check presets exist + if 'presets' in filters: + presets = filters['presets'] + log_test( + "Filter presets defined", + len(presets) > 0, + f"Presets: {list(presets.keys())}" + ) + else: + log_test("filters.yaml exists", False, "File not found") + + # Test analyses.yaml + analyses_file = Path("data/config/analyses.yaml") + if analyses_file.exists(): + with open(analyses_file) as f: + analyses = yaml.safe_load(f) + + log_test( + "analyses.yaml exists and valid", + isinstance(analyses, dict) and 'analyses' in analyses, + f"Keys: {list(analyses.keys())}" + ) + + # Check built-in analyses documented + if 'analyses' in analyses: + builtin = analyses['analyses'] + expected_analyses = ['xrd_dara', 'powder_statistics'] + for analyzer in expected_analyses: + log_test( + f"Analysis '{analyzer}' documented", + analyzer in builtin, + f"Documented: {list(builtin.keys())}" + ) + else: + log_test("analyses.yaml exists", False, "File not found") + + return True + + except Exception as e: + log_test("Configuration files", False, str(e)) + traceback.print_exc() + return False + + +def test_mongodb_connection(): + """Test MongoDB connection (optional)""" + test_header("TEST 8: MongoDB Connection (Optional)") + + try: + from pymongo import MongoClient + + client = MongoClient("mongodb://localhost:27017/", serverSelectionTimeoutMS=2000) + db = client["temporary"] + collection = db["release"] + + # Try to count documents + count = collection.count_documents({}) + log_test( + "MongoDB connection successful", + count >= 0, + f"Found {count} experiments in database" + ) + + # Test aggregation pipeline + try: + pipeline = [ + {"$group": {"_id": {"$substr": ["$name", 0, 3]}, "count": {"$sum": 1}}}, + {"$sort": {"count": -1}}, + {"$limit": 5} + ] + results = list(collection.aggregate(pipeline)) + log_test( + "MongoDB aggregation works", + len(results) > 0, + f"Top types: {[r['_id'] for r in results]}" + ) + except Exception as e: + log_test("MongoDB aggregation", False, str(e)) + + client.close() + return True + + except Exception as e: + log_test("MongoDB connection", False, f"Not available: {e}") + print(" ℹ MongoDB tests skipped (server not available)") + return True # Don't fail if MongoDB isn't running + + +def test_parquet_transformer(): + """Test parquet transformation with filters""" + test_header("TEST 9: Parquet Transformer with Filters") + + try: + from mongodb_to_parquet import MongoToParquetTransformer + + # Create a temp directory for test output + with tempfile.TemporaryDirectory() as tmpdir: + output_dir = Path(tmpdir) / "test_parquet" + + # Test with filter (limit to 5 experiments for speed) + experiment_filter = { + 'status': ['completed'], + 'has_xrd': True + } + + try: + transformer = MongoToParquetTransformer(output_dir=output_dir) + # Use limit to make test fast + transformer.transform_all( + limit=5, + skip_temp_logs=True, + skip_xrd_points=True, + experiment_filter=experiment_filter + ) + transformer.close() + + # Check output files - metadata.json should always exist + metadata_file = output_dir / 'metadata.json' + log_test( + "Generated metadata.json", + metadata_file.exists(), + f"Size: {metadata_file.stat().st_size if metadata_file.exists() else 0} bytes" + ) + + # Parquet files may be empty if no data matches filter + # This is OK - just log for informational purposes + if (output_dir / 'experiments.parquet').exists(): + size = (output_dir / 'experiments.parquet').stat().st_size + log_test( + "Generated experiments.parquet", + True, + f"Size: {size} bytes (0 bytes = no matching data, OK)" + ) + else: + log_test("experiments.parquet not created", True, "No matching data (OK)") + + log_test("Parquet transformation with filters", True, "Successfully generated filtered parquet") + + except Exception as e: + log_test("Parquet transformation", False, f"MongoDB not available or error: {e}") + + return True + + except Exception as e: + log_test("Parquet transformer", False, str(e)) + print(" ℹ Transformer test skipped (MongoDB may not be available)") + return True + + +def test_diagram_generation(): + """Test schema diagram generation""" + test_header("TEST 10: Diagram Generation") + + try: + # Check if we have parquet data to diagram + parquet_dir = Path("data/parquet") + + if not parquet_dir.exists() or not list(parquet_dir.glob("*.parquet")): + log_test( + "Diagram generation", + True, + "Skipped: No parquet data available" + ) + return True + + from generate_diagram import ParquetSchemaAnalyzer, DiagramGenerator + + # Analyze schema + analyzer = ParquetSchemaAnalyzer(parquet_dir) + analysis = analyzer.analyze() + + log_test( + "Schema analysis", + 'tables' in analysis and len(analysis['tables']) > 0, + f"Found {len(analysis['tables'])} tables" + ) + + # Generate diagram + generator = DiagramGenerator(analysis) + + # Test terminal output (shouldn't error) + try: + mermaid = generator.generate_mermaid() + log_test( + "Mermaid ERD generation", + '```mermaid' in mermaid and 'erDiagram' in mermaid, + f"Generated {len(mermaid)} characters" + ) + except Exception as e: + log_test("Mermaid generation", False, str(e)) + + # Test summary generation + try: + summary = generator.generate_summary() + log_test( + "Summary generation", + '# Parquet Schema Documentation' in summary, + f"Generated {len(summary)} characters" + ) + except Exception as e: + log_test("Summary generation", False, str(e)) + + return True + + except Exception as e: + log_test("Diagram generation", False, str(e)) + traceback.print_exc() + return False + + +def test_full_integration(): + """Test full pipeline integration""" + test_header("TEST 11: Full Pipeline Integration") + + try: + # Check all required directories exist + dirs_to_check = [ + Path("data/products"), + Path("data/products/schema"), + Path("data/pipeline"), + Path("data/analyses"), + Path("data/config"), + Path("data"), + Path("data/tools") + ] + + all_exist = True + for dir_path in dirs_to_check: + exists = dir_path.exists() + log_test(f"Directory {dir_path} exists", exists) + if not exists: + all_exist = False + + # Check core files + files_to_check = [ + Path("data/products/schema_manager.py"), + Path("data/products/schema_validator.py"), + Path("data/analyses/base_analyzer.py"), + Path("data/pipeline/product_pipeline.py"), + Path("data/config/defaults.yaml"), + Path("data/config/filters.yaml"), + Path("data/config/analyses.yaml"), + Path("run_product_pipeline.sh") + ] + + for file_path in files_to_check: + exists = file_path.exists() + log_test(f"File {file_path.name} exists", exists) + if not exists: + all_exist = False + + # Check Python dependencies + try: + import pandas + import pydantic + import yaml + import rich + log_test("Core Python dependencies installed", True) + except ImportError as e: + log_test("Core Python dependencies", False, str(e)) + all_exist = False + + return all_exist + + except Exception as e: + log_test("Full integration", False, str(e)) + traceback.print_exc() + return False + + +# ============================================================================= +# Main Test Runner +# ============================================================================= + +def main(): + """Run all integration tests""" + print("\n" + "=" * 70) + print(" A-LAB PIPELINE COMPREHENSIVE INTEGRATION TESTS") + print("=" * 70) + print(f"\n Started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") + + # Run all tests + tests = [ + test_schema_auto_discovery, + test_analysis_auto_discovery, + test_custom_schema_hook, + test_custom_analyzer_hook, + test_filter_configurations, + test_edge_cases, + test_config_files, + test_mongodb_connection, + test_parquet_transformer, + test_diagram_generation, + test_full_integration + ] + + total_tests = 0 + passed_tests = 0 + + for test_func in tests: + try: + result = test_func() + # Count individual test results logged within each function + except Exception as e: + print(f"\n✗ Test suite {test_func.__name__} crashed: {e}") + traceback.print_exc() + + # Summary + print("\n" + "=" * 70) + print(" TEST SUMMARY") + print("=" * 70) + + passed = sum(1 for r in test_results if r['passed']) + failed = sum(1 for r in test_results if not r['passed']) + total = len(test_results) + + print(f"\n Total: {total} tests") + print(f" Passed: {passed} tests ({'green' if passed == total else 'yellow'})") + print(f" Failed: {failed} tests ({'red' if failed > 0 else 'green'})") + + if failed > 0: + print("\n Failed tests:") + for result in test_results: + if not result['passed']: + print(f" ✗ {result['test']}") + if result['message']: + print(f" {result['message']}") + + print("\n" + "=" * 70) + + if failed == 0: + print(" ✓ ALL TESTS PASSED!") + print("\n Next steps:") + print(" 1. Create a product: ./run_product_pipeline.sh create") + print(" 2. Run pipeline: ./run_product_pipeline.sh run --product ") + print("=" * 70) + return 0 + else: + print(" ✗ SOME TESTS FAILED") + print("\n Review failures above and fix issues.") + print("=" * 70) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/SCHEMA_SYSTEM.md b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/SCHEMA_SYSTEM.md new file mode 100644 index 000000000..1b1530656 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/SCHEMA_SYSTEM.md @@ -0,0 +1,202 @@ +# A-Lab Schema System (Pydantic-First) + +This document describes the Pydantic-first schema system for A-Lab data validation and MPContribs uploads. + +## Overview + +**Schemas are defined directly in Python using Pydantic** - not generated from YAML. This approach: + +- Aligns with the team's existing `results_schema.py` +- Enables rich validation (constraints, validators, Literal types) +- Better IDE support (autocomplete, type checking) +- Matches MPContribs expectations + +## Schema Files + +All schemas are in `data/products/schema/`: + +``` +data/products/schema/ +├── __init__.py # Package exports +├── base.py # ExcludeFromUpload utility +├── experiments.py # Main consolidated table (~37 fields) +├── experiment_elements.py # Elements per experiment (1:N) +├── powder_doses.py # Individual powder doses (1:N) +├── temperature_logs.py # Temperature readings (1:N, optional) +├── xrd_data_points.py # Raw XRD patterns (1:N, optional) +├── workflow_tasks.py # Task execution history (1:N) +├── xrd_refinements.py # DARA analysis results +└── xrd_phases.py # Identified crystal phases +``` + +**One schema per parquet file** - no redundancy. +All constraints from team's `results_schema.py` are integrated into these schemas. + +## Parquet Table Schemas + +Each schema corresponds to **one parquet file**. All constraints from team's `results_schema.py` are integrated: + +| Schema | Class | Table | Relationship | Constraints | +| ------------------------ | --------------------- | --------------------------- | ------------ | ------------------------ | +| `experiments.py` | `Experiment` | experiments.parquet | Main table | Literal types, ranges | +| `experiment_elements.py` | `ExperimentElement` | experiment_elements.parquet | 1:N | - | +| `powder_doses.py` | `PowderDose` | powder_doses.parquet | 1:N | ge=0 for masses | +| `temperature_logs.py` | `TemperatureLogEntry` | temperature_logs.parquet | 1:N | - | +| `workflow_tasks.py` | `WorkflowTask` | workflow_tasks.parquet | 1:N | Literal status | +| `xrd_data_points.py` | `XRDDataPoint` | xrd_data_points.parquet | 1:N | ge=0 for point_index | +| `xrd_refinements.py` | `XRDRefinement` | xrd_refinements.parquet | 1:1 | - | +| `xrd_phases.py` | `XRDPhase` | xrd_phases.parquet | 1:N | 0 <= weight_fraction <=1 | + +### Data Flow + +MongoDB (nested) → Flattened → Parquet (normalized) → Schemas validate parquet + +The team's `results_schema.py` describes MongoDB structure. Our schemas describe the **flattened parquet** representation of that same data. + +## Field Exclusion (Embargoed Data) + +Some fields should NOT be uploaded to MPContribs until the paper is published: + +```python +from .base import ExcludeFromUpload + +class RecoverPowderResult(BaseModel, extra="forbid"): + # EXCLUDED FROM UPLOAD - embargoed until paper publication + weight_collected: float | None = ExcludeFromUpload( + description="Weight of powder collected (EMBARGOED)" + ) +``` + +**Currently excluded fields:** + +- `recovery_weight_collected_mg` - powder recovery weight +- `xrd_total_mass_dispensed_mg` - XRD mass dispensed +- `mass_per_dispensing_attempt_mg` - per-attempt XRD mass +- `first_tapping_mass_collected` - first tapping mass + +## Using Schemas + +### SchemaManager + +```python +from data.products.schema_manager import SchemaManager + +sm = SchemaManager() + +# List all tables +print(sm.get_table_names()) +# ['experiments', 'experiment_elements', 'powder_doses', ...] + +# Get schema class +Experiment = sm.get_schema('experiments') + +# Get fields for upload (excluding embargoed) +uploadable = sm.get_uploadable_fields('experiments') +# ['experiment_id', 'name', ...] (35 of 37 fields) + +excluded = sm.get_excluded_fields('experiments') +# ['recovery_weight_collected_mg', 'xrd_total_mass_dispensed_mg'] +``` + +### Validation + +```python +from data.products.schema_validator import SchemaValidator + +validator = SchemaValidator() +is_valid, errors, warnings = validator.validate_parquet_data(parquet_dir) +``` + +### Direct Import + +```python +from data.products.schema import ( + Experiment, + PowderDose, + XRDRefinement, + XRDPhase, +) + +# Validate a row +exp = Experiment( + experiment_id="abc123", + name="NSC_249", + experiment_type="NSC", + target_formula="Na3V2(PO4)3", + status="completed", + last_updated=datetime.now(), + # Heating data (flattened from MongoDB metadata.heating_results) + heating_temperature=800.0, + heating_time=240.0, + # Recovery data (flattened from MongoDB metadata.recoverpowder_results) + recovery_initial_crucible_weight_mg=1500.0, + # Note: recovery_weight_collected_mg is EXCLUDED from upload (embargoed) +) +``` + +## Upload to S3 OpenData + +The pipeline now uploads parquet files directly to S3: + +```python +from data.pipeline.s3_uploader import upload_product_to_s3 + +# Upload (dry run) +uploaded = upload_product_to_s3( + product_name="my_product", + product_dir=Path("data/products/my_product"), + dry_run=True +) + +# Actual upload +uploaded = upload_product_to_s3(..., dry_run=False) +``` + +Files are uploaded to: `s3://materialsproject-contribs/alab_synthesis/{product_name}/` + +## Modifying Schemas + +To add or modify fields, **edit the Python files directly**: + +1. Open `data/products/schema/experiments.py` +2. Add/modify fields with Pydantic syntax +3. Run pipeline - validation will use updated schema automatically + +```python +# Add a new field +new_field: str | None = Field( + default=None, + description="Description of the new field" +) + +# Mark a field as excluded from upload +sensitive_field: float | None = ExcludeFromUpload( + description="Sensitive data (EMBARGOED)" +) + +# Add constraints +position: int | None = Field( + default=None, + description="Position index", + ge=1, # >= 1 + le=16 # <= 16 +) + +# Use Literal for enum-like fields +status: Literal["completed", "error", "active", "unknown"] = Field(...) +``` + +## Parquet-Schema Alignment + +The `parquet_data_loader.py` dashboard loader is aligned with the schemas: + +| Schema Field | Parquet Column | +| ------------------------------ | ------------------------------ | +| `heating_temperature` | `heating_temperature` | +| `heating_time` | `heating_time` | +| `heating_cooling_rate` | `heating_cooling_rate` | +| `recovery_weight_collected_mg` | `recovery_weight_collected_mg` | +| `xrd_sampleid_in_aeris` | `xrd_sampleid_in_aeris` | +| `dosing_crucible_position` | `dosing_crucible_position` | + +All prefixed fields (heating*, recovery*, xrd*, dosing*, finalization\_) are consistent between schemas and parquet files. diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/base_product.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/base_product.py new file mode 100644 index 000000000..c4f0f87e5 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/base_product.py @@ -0,0 +1,771 @@ +#!/usr/bin/env python3 +""" +Base Product Configuration System for A-Lab Data Products + +Provides a flexible framework for defining data products with: +- Experiment filtering (by type, status, date, etc.) +- Schema definition with units for MPContribs +- Pydantic validation +- Analysis pipeline configuration +- Metadata for publications +""" + +import yaml +import json +import sys +from pathlib import Path +from datetime import datetime +from typing import Dict, List, Optional, Any, Type +from pydantic import BaseModel, Field, validator, create_model +from enum import Enum +import pandas as pd +import subprocess +from rich.console import Console + +# Import config loader +sys.path.insert(0, str(Path(__file__).parent.parent / "config")) +from config_loader import get_config + +console = Console() + + +class ExperimentType(str, Enum): + """Known experiment type prefixes""" + NSC = "NSC" + Na = "Na" + PG = "PG" + MINES = "MINES" + TRI = "TRI" + + +class ExperimentStatus(str, Enum): + """Experiment workflow status""" + COMPLETED = "completed" + ERROR = "error" + ACTIVE = "active" + UNKNOWN = "unknown" + + +class AnalysisConfig(BaseModel): + """Configuration for a single analysis module""" + name: str + enabled: bool = True + config: Dict[str, Any] = {} + + +class SchemaField(BaseModel): + """Schema field definition for MPContribs""" + unit: Optional[str] = Field(None, description="Unit for numeric fields, '' for strings, None for dimensionless") + description: str = Field(..., description="Human-readable description") + type: str = Field("string", description="Data type: float, int, boolean, string") + required: bool = Field(False, description="Whether field is required") + min: Optional[float] = None + max: Optional[float] = None + + +class ExperimentFilter(BaseModel): + """Filter criteria for selecting experiments""" + types: Optional[List[ExperimentType]] = Field(None, description="Experiment type prefixes") + status: Optional[List[ExperimentStatus]] = Field(None, description="Workflow status") + has_xrd: Optional[bool] = Field(None, description="Must have XRD data") + has_sem: Optional[bool] = Field(None, description="Must have SEM data") + date_range: Optional[Dict[str, str]] = Field(None, description="Start/end dates") + experiment_names: Optional[List[str]] = Field(None, description="Specific experiment names") + + def to_mongo_query(self) -> Dict: + """Convert filter to MongoDB query""" + query = {} + + # Handle name filtering (types OR specific experiment names) + name_queries = [] + + if self.types: + # Match experiments starting with any of the prefixes + # Convert enum values to strings + type_strings = [t.value if hasattr(t, 'value') else str(t) for t in self.types] + patterns = [f"^{t}_" for t in type_strings] + name_queries.append({"name": {"$regex": "|".join(patterns)}}) + + if self.experiment_names: + # Specific experiment names + name_queries.append({"name": {"$in": self.experiment_names}}) + + # Combine name queries with $or if both exist + if len(name_queries) == 1: + query.update(name_queries[0]) + elif len(name_queries) > 1: + query["$or"] = name_queries + + if self.status: + query["status"] = {"$in": [s.value for s in self.status]} + + if self.has_xrd is not None: + if self.has_xrd: + query["metadata.diffraction_results.sampleid_in_aeris"] = {"$exists": True, "$ne": None} + else: + # Need to handle $or carefully if it already exists + xrd_or = [ + {"metadata.diffraction_results.sampleid_in_aeris": {"$exists": False}}, + {"metadata.diffraction_results.sampleid_in_aeris": None} + ] + if "$or" in query: + # Wrap both in $and + existing_or = query.pop("$or") + query["$and"] = [ + {"$or": existing_or}, + {"$or": xrd_or} + ] + else: + query["$or"] = xrd_or + + if self.has_sem is not None: + # Check for SEM data existence + if self.has_sem: + query["metadata.sem_results"] = {"$exists": True, "$ne": None} + + if self.date_range: + date_query = {} + if "start" in self.date_range: + date_query["$gte"] = datetime.fromisoformat(self.date_range["start"]) + if "end" in self.date_range: + date_query["$lte"] = datetime.fromisoformat(self.date_range["end"]) + if date_query: + query["last_updated"] = date_query + + return query + + +class ProductMetadata(BaseModel): + """Metadata for MPContribs project""" + title: Optional[str] = Field(None, description="Human-readable title") + authors: Optional[str] = Field(None, description="Author list") + description: Optional[str] = Field(None, description="Abstract/description") + references: Optional[List[Dict[str, str]]] = Field(None, description="List of {label, url}") + doi: Optional[str] = Field(None, description="DOI if published") + + +class ProductConfig(BaseModel): + """Complete configuration for a data product""" + name: str = Field(..., description="Product identifier (alphanumeric + underscore)") + version: str = Field("1.0", description="Product version") + schema_version: str = Field("1.0", description="Schema version (for migration tracking)") + + experiment_filter: ExperimentFilter + metadata: ProductMetadata = ProductMetadata() + analyses: List[AnalysisConfig] = [] + data_schema: Dict[str, SchemaField] = Field(default_factory=dict, alias='schema') + + created_at: datetime = Field(default_factory=datetime.now) + updated_at: datetime = Field(default_factory=datetime.now) + submitted_to_mpcontribs: bool = Field(False, description="Whether product has been submitted to MPContribs") + + model_config = { + 'populate_by_name': True # Allow both 'data_schema' and 'schema' to work + } + + @validator("name") + def validate_name(cls, v): + """Ensure name is valid identifier""" + if not v.replace("_", "").isalnum(): + raise ValueError("Product name must be alphanumeric with underscores only") + return v + + def save(self, path: Path): + """Save configuration to YAML file""" + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + + # Convert to dict with serializable types + config_dict = json.loads(self.json()) + + with open(path, "w") as f: + yaml.dump(config_dict, f, default_flow_style=False, sort_keys=False) + + @classmethod + def load(cls, path: Path) -> "ProductConfig": + """Load configuration from YAML file""" + with open(path, "r") as f: + data = yaml.safe_load(f) + return cls(**data) + + def to_pydantic_model(self) -> Type[BaseModel]: + """ + Generate Pydantic model from schema definition + + Returns a dynamically created Pydantic model class for validation + """ + fields = {} + validators = {} + + for field_name, field_def in self.data_schema.items(): + # Determine Python type + if field_def.type == "float": + python_type = float + elif field_def.type == "int": + python_type = int + elif field_def.type == "boolean": + python_type = bool + else: + python_type = str + + # Make optional if not required + if not field_def.required: + python_type = Optional[python_type] + + # Create field with description + fields[field_name] = (python_type, Field(description=field_def.description)) + + # Add validators for min/max + if field_def.min is not None or field_def.max is not None: + def make_validator(fname, fmin, fmax): + def validator(cls, v): + if v is not None: + if fmin is not None and v < fmin: + raise ValueError(f"{fname} must be >= {fmin}") + if fmax is not None and v > fmax: + raise ValueError(f"{fname} must be <= {fmax}") + return v + return validator + + validators[f"validate_{field_name}"] = validator(field_name)( + make_validator(field_name, field_def.min, field_def.max) + ) + + # Create dynamic model + model_name = f"{self.name.title().replace('_', '')}Experiment" + return create_model(model_name, **fields, __validators__=validators) + + def get_mpcontribs_columns(self) -> Dict[str, Optional[str]]: + """ + Get columns with units for MPContribs initialization + + Returns: + Dict of column_name -> unit (None for dimensionless, "" for strings) + """ + return { + self._to_camel_case(name): field.unit + for name, field in self.data_schema.items() + } + + @staticmethod + def _to_camel_case(snake_str: str) -> str: + """Convert snake_case to camelCase for MPContribs""" + components = snake_str.split('_') + return components[0] + ''.join(x.title() for x in components[1:]) + + +class ProductManager: + """Manager for data product lifecycle""" + + def __init__(self, products_dir: Path = None): + self.products_dir = Path(products_dir or "data/products") + self.products_dir.mkdir(parents=True, exist_ok=True) + + def create_product_interactive(self) -> ProductConfig: + """Interactive CLI for creating a new product configuration""" + import inquirer + from rich.console import Console + from rich.table import Table + from rich.tree import Tree + + console = Console() + + console.print("\n[bold cyan]Create New Data Product[/bold cyan]\n") + + # Product name + name = inquirer.text( + message="Product name (alphanumeric + underscore)", + validate=lambda _, x: x.replace("_", "").isalnum() + ) + + # Discover experiment types using MongoDB + console.print("\n[yellow]Discovering experiment types from MongoDB...[/yellow]") + hierarchy = self._discover_experiment_types() + + # Display hierarchical structure + tree = Tree("[bold cyan]Available Experiment Groups[/bold cyan]") + + for root, data in sorted(hierarchy.items()): + root_node = tree.add(f"[green]{root}[/green] ({data['count']} experiments)") + + if data['subgroups']: + for subgroup, subdata in sorted(data['subgroups'].items()): + root_node.add(f"[yellow]{subgroup}[/yellow] ({subdata['count']} experiments)") + + console.print(tree) + console.print() + + # Calculate total experiments + total_experiments = sum(data['count'] for data in hierarchy.values()) + + console.print(f"[yellow]ℹ Selection behavior:[/yellow]") + console.print(f" • Select specific groups to filter experiments") + console.print(f" • Leave empty to include ALL {total_experiments} experiments (will require confirmation)") + console.print() + + # Build selection choices with hierarchy + all_choices = [] + choice_map = {} # Maps display string to selection value + + for root in sorted(hierarchy.keys()): + data = hierarchy[root] + + # Add root level option + root_label = f"{root} (all {data['count']} experiments)" + all_choices.append(root_label) + choice_map[root_label] = {'type': 'root', 'value': root} + + # Add subgroup options + if data['subgroups']: + for subgroup in sorted(data['subgroups'].keys()): + subdata = data['subgroups'][subgroup] + subgroup_label = f" └─ {subgroup} ({subdata['count']} experiments)" + all_choices.append(subgroup_label) + choice_map[subgroup_label] = {'type': 'subgroup', 'value': subgroup} + + # Select experiment groups + selected_labels = inquirer.checkbox( + message="Select experiment groups (leave empty for ALL experiments)", + choices=all_choices + ) + + # If nothing selected, confirm they want ALL experiments + if not selected_labels: + console.print(f"\n[yellow]⚠ No groups selected - this will include ALL {total_experiments} experiments![/yellow]") + confirm_all = inquirer.confirm( + message=f"Proceed with ALL {total_experiments} experiments?", + default=False + ) + + if not confirm_all: + console.print("\n[cyan]Please select specific experiment groups...[/cyan]\n") + selected_labels = inquirer.checkbox( + message="Select experiment groups", + choices=all_choices + ) + + # If still nothing selected, abort + if not selected_labels: + console.print("\n[red]No experiments selected. Aborting product creation.[/red]") + return None + + # Process selections into types and experiment_names + selected_roots = set() + selected_subgroups = [] + selected_experiment_names = [] + + for label in selected_labels: + choice = choice_map[label] + if choice['type'] == 'root': + selected_roots.add(choice['value']) + elif choice['type'] == 'subgroup': + selected_subgroups.append(choice['value']) + # Get all experiments for this subgroup + root = choice['value'].split('_')[0] + if root in hierarchy and choice['value'] in hierarchy[root]['subgroups']: + experiments = hierarchy[root]['subgroups'][choice['value']]['experiments'] + selected_experiment_names.extend(experiments) + + # Build filter config + filter_config = {} + + # If nothing selected (user confirmed to use ALL), don't add type/name filters + # This will match all experiments in the database + if not selected_roots and not selected_experiment_names: + console.print(f"[dim]Filter: ALL experiments (no type/name filter applied)[/dim]") + # If only subgroups selected (no full roots), use experiment_names + elif selected_experiment_names and not selected_roots: + filter_config["experiment_names"] = selected_experiment_names + console.print(f"[dim]Filter: {len(selected_experiment_names)} specific experiments[/dim]") + # If roots selected, use types (and ignore subgroups under those roots) + elif selected_roots: + filter_config["types"] = list(selected_roots) + console.print(f"[dim]Filter: Types {list(selected_roots)}[/dim]") + # If both, combine: use types for roots, add experiment_names for other subgroups + elif selected_roots and selected_subgroups: + filter_config["types"] = list(selected_roots) + # Only add experiment names from subgroups whose root isn't selected + filtered_names = [ + exp for exp in selected_experiment_names + if exp.split('_')[0] not in selected_roots + ] + if filtered_names: + filter_config["experiment_names"] = filtered_names + console.print(f"[dim]Filter: Types {list(selected_roots)} + {len(filtered_names)} specific experiments[/dim]") + + # Additional filters - let user select which ones to apply + console.print("\n[cyan]Additional Filters[/cyan]") + available_filters = inquirer.checkbox( + message="Select additional filters to apply (optional)", + choices=[ + "XRD requirement", + "Status filter", + "SEM requirement", + "Date range" + ] + ) + + # Configure selected filters + if "XRD requirement" in available_filters: + xrd_options = inquirer.checkbox( + message="XRD requirement (select one)", + choices=["Must have XRD", "Must NOT have XRD"], + carousel=True + ) + if "Must have XRD" in xrd_options: + filter_config["has_xrd"] = True + elif "Must NOT have XRD" in xrd_options: + filter_config["has_xrd"] = False + + if "Status filter" in available_filters: + status = inquirer.checkbox( + message="Select experiment statuses to include", + choices=["completed", "error", "active", "unknown"], + default=["completed"] + ) + if status: + filter_config["status"] = status + + if "SEM requirement" in available_filters: + sem_options = inquirer.checkbox( + message="SEM requirement (select one)", + choices=["Must have SEM", "Must NOT have SEM"], + carousel=True + ) + if "Must have SEM" in sem_options: + filter_config["has_sem"] = True + elif "Must NOT have SEM" in sem_options: + filter_config["has_sem"] = False + + if "Date range" in available_filters: + start_date = inquirer.text( + message="Start date (YYYY-MM-DD, or blank for no start)", + default="" + ) + end_date = inquirer.text( + message="End date (YYYY-MM-DD, or blank for no end)", + default="" + ) + date_range = {} + if start_date: + date_range["start"] = start_date + if end_date: + date_range["end"] = end_date + if date_range: + filter_config["date_range"] = date_range + + # Metadata (optional) + include_metadata = inquirer.confirm( + message="Include publication metadata?", + default=False + ) + + metadata = {} + if include_metadata: + metadata["title"] = inquirer.text(message="Title") + metadata["authors"] = inquirer.text(message="Authors") + metadata["description"] = inquirer.text(message="Description") + + # Analyses + analyses = [] + available_analyses = self._discover_analyses() + + if available_analyses: + console.print("\n[cyan]Available Analyses:[/cyan]") + for analysis in available_analyses: + console.print(f" • {analysis}") + + selected_analyses = inquirer.checkbox( + message="Select analyses to run", + choices=available_analyses + ) + + analyses = [{"name": a, "enabled": True} for a in selected_analyses] + + # Show schema info (schemas are managed centrally in schema/ directory) + console.print("\n[cyan]A-Lab Pydantic Schemas[/cyan]") + from schema_manager import SchemaManager + + schema_manager = SchemaManager() + experiments_schema = schema_manager.get_main_schema() + + if experiments_schema: + # Get field info from Pydantic model + field_count = len(experiments_schema.model_fields) + excluded_fields = schema_manager.get_excluded_fields('experiments') + + console.print(f"[green]✓ Experiments schema: {field_count} fields[/green]") + if excluded_fields: + console.print(f" [yellow]⚠ {len(excluded_fields)} fields excluded from upload (embargoed)[/yellow]") + + # Show summary of fields by prefix + console.print("\n[dim]Schema includes:[/dim]") + all_fields = list(experiments_schema.model_fields.keys()) + prefixes = { + 'Core': [f for f in all_fields if f in ['experiment_id', 'name', 'experiment_type', 'target_formula', 'status']], + 'Heating': [f for f in all_fields if f.startswith('heating_')], + 'Recovery': [f for f in all_fields if f.startswith('recovery_')], + 'XRD': [f for f in all_fields if f.startswith('xrd_')], + 'Dosing': [f for f in all_fields if f.startswith('dosing_')], + 'Finalization': [f for f in all_fields if f.startswith('finalization_')] + } + + for category, fields in prefixes.items(): + if fields: + console.print(f" [yellow]{category}[/yellow]: {len(fields)} fields") + + console.print() + console.print("[dim]Note: Schemas are managed centrally in data/products/schema/[/dim]") + console.print("[dim]Edit the .py files directly to modify validation rules[/dim]") + else: + console.print(f"[yellow]⚠ Experiments schema not found[/yellow]") + + # Create config (no data_schema - managed centrally now) + config = ProductConfig( + name=name, + experiment_filter=ExperimentFilter(**filter_config), + metadata=ProductMetadata(**metadata) if metadata else ProductMetadata(), + analyses=analyses, + data_schema={} # Empty - schemas are managed centrally in schema/ directory + ) + + # Save + config_path = self.products_dir / name / "config.yaml" + config.save(config_path) + + console.print(f"\n[green]✓ Created configuration: {config_path}[/green]") + + # Show Pydantic schemas info (schemas are Python-first, not generated) + console.print(f"\n[cyan]Pydantic Schemas (from data/products/schema/):[/cyan]") + for table_name in schema_manager.get_table_names(): + schema = schema_manager.get_schema(table_name) + excluded = schema_manager.get_excluded_fields(table_name) + excluded_str = f" (excluded: {', '.join(excluded)})" if excluded else "" + console.print(f" • {table_name}: {len(schema.model_fields)} fields{excluded_str}") + + return config + + def _discover_experiment_types(self) -> Dict[str, Any]: + """ + Discover experiment types with hierarchical grouping based on underscores + + Returns: + Dict with hierarchical structure: + { + 'NSC': { + 'count': 218, + 'subgroups': { + 'NSC_249': {'count': 5, 'experiments': ['NSC_249_001', ...]}, + 'NSC_250': {'count': 3, 'experiments': [...]} + } + }, + ... + } + """ + from pymongo import MongoClient + from collections import defaultdict + + try: + # Connect to MongoDB (using config loader) + config = get_config() + client = MongoClient(config.mongo_uri, serverSelectionTimeoutMS=5000) + db = client[config.mongo_db] + collection = db[config.mongo_collection] + + # Get all experiment names + experiments = collection.find({}, {"name": 1, "_id": 0}) + + # Build hierarchical structure + hierarchy = defaultdict(lambda: { + 'count': 0, + 'subgroups': defaultdict(lambda: { + 'count': 0, + 'experiments': [] + }) + }) + + for exp in experiments: + exp_name = exp.get('name', '') + if not exp_name or '_' not in exp_name: + continue + + # Split by underscore to get all levels + parts = exp_name.split('_') + + # Root level (e.g., 'NSC', 'Na') + root = parts[0] + hierarchy[root]['count'] += 1 + + # If there are multiple parts, create subgroups + if len(parts) >= 2: + # Build all possible subgroup prefixes + # For NSC_249_001, create: NSC_249 + # For Na_123_A_001, create: Na_123, Na_123_A + for i in range(1, len(parts)): + subgroup = '_'.join(parts[:i+1]) + + # Only track subgroups that aren't the full experiment name + # (i.e., there's at least one more part after this) + if i < len(parts) - 1: + hierarchy[root]['subgroups'][subgroup]['count'] += 1 + hierarchy[root]['subgroups'][subgroup]['experiments'].append(exp_name) + + client.close() + + # Convert defaultdict to regular dict for serialization + result = {} + for root, data in hierarchy.items(): + result[root] = { + 'count': data['count'], + 'subgroups': dict(data['subgroups']) if data['subgroups'] else {} + } + + return result + + except Exception as e: + console.print(f"[yellow]Warning: Could not connect to MongoDB: {e}[/yellow]") + console.print("[yellow]Using fallback experiment types[/yellow]") + # Return empty hierarchy as fallback + return { + "NSC": {"count": 0, "subgroups": {}}, + "Na": {"count": 0, "subgroups": {}}, + "PG": {"count": 0, "subgroups": {}}, + "MINES": {"count": 0, "subgroups": {}}, + "TRI": {"count": 0, "subgroups": {}} + } + + def _discover_analyses(self) -> List[str]: + """Discover available analysis modules""" + analyses = ["xrd_dara"] # Always available + + # Check for other analysis scripts + analysis_paths = [ + Path("data/analyses/powder_statistics.py"), + Path("data/analyses/sem_clustering.py"), + Path("data/analyses/heating_profile.py") + ] + + for path in analysis_paths: + if path.exists(): + analyses.append(path.stem) + + return analyses + + def _generate_pydantic_schema(self, config: ProductConfig): + """ + Generate Pydantic validation schema file from product config + + This is auto-generated from the YAML schema and should be regenerated + whenever the schema changes. + """ + schema_path = self.products_dir / config.name / "schema.py" + + # Generate Python code for the schema + code = f'''#!/usr/bin/env python3 +""" +Auto-generated Pydantic schema for {config.name} +Generated: {datetime.now().isoformat()} + +⚠️ DO NOT EDIT THIS FILE MANUALLY ⚠️ +This file is auto-generated from config.yaml +To modify the schema, edit config.yaml and regenerate with: + python data/products/base_product.py regenerate {config.name} +""" + +from pydantic import BaseModel, Field, validator +from typing import Optional +from datetime import datetime as dt_type + + +class {config.name.title().replace("_", "")}Experiment(BaseModel): + """ + Validation schema for {config.name} experiments + + This schema validates data before upload to MPContribs. + All fields are derived from the product's data_schema in config.yaml. + """ + + # Auto-generated fields +''' + + # Add fields + for field_name, field_def in config.data_schema.items(): + python_type = { + "float": "float", + "int": "int", + "boolean": "bool", + "string": "str", + "datetime": "dt_type" + }.get(field_def.type, "str") + + if not field_def.required: + python_type = f"Optional[{python_type}]" + + default = " = None" if not field_def.required else "" + + # Add unit info to docstring if present + unit_str = f" ({field_def.unit})" if field_def.unit else "" + code += f" {field_name}: {python_type}{default} # {field_def.description}{unit_str}\n" + + # Add validators + validators_added = False + for field_name, field_def in config.data_schema.items(): + if field_def.min is not None or field_def.max is not None: + if not validators_added: + code += '\n # Validators for numeric ranges\n' + validators_added = True + + code += f''' + @validator("{field_name}") + def validate_{field_name}(cls, v): + """Validate {field_name} is within acceptable range""" + if v is not None: +''' + if field_def.min is not None: + code += f' if v < {field_def.min}:\n' + code += f' raise ValueError("{field_name} must be >= {field_def.min}")\n' + if field_def.max is not None: + code += f' if v > {field_def.max}:\n' + code += f' raise ValueError("{field_name} must be <= {field_def.max}")\n' + code += ' return v\n' + + # Add metadata + code += f''' + +# Metadata +__schema_version__ = "{config.version}" +__generated_at__ = "{datetime.now().isoformat()}" +__source_file__ = "config.yaml" +__field_count__ = {len(config.data_schema)} +''' + + # Write schema file + with open(schema_path, 'w') as f: + f.write(code) + + console.print(f"[green]✓ Generated Pydantic schema: {schema_path}[/green]") + console.print(f"[dim] {len(config.data_schema)} fields, {sum(1 for f in config.data_schema.values() if f.required)} required[/dim]") + + def list_products(self) -> List[str]: + """List all available products""" + products = [] + for path in self.products_dir.iterdir(): + if path.is_dir() and (path / "config.yaml").exists(): + products.append(path.name) + return products + + def get_product_config(self, name: str) -> ProductConfig: + """Load product configuration by name""" + config_path = self.products_dir / name / "config.yaml" + if not config_path.exists(): + raise FileNotFoundError(f"Product '{name}' not found") + return ProductConfig.load(config_path) + + +if __name__ == "__main__": + # Test the interactive product creation + manager = ProductManager() + config = manager.create_product_interactive() + print(f"\nCreated product: {config.name}") + print(f"Experiments filter: {config.experiment_filter}") + print(f"Schema fields: {list(config.data_schema.keys())}") diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/schema_manager.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/schema_manager.py new file mode 100644 index 000000000..857cb5321 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/schema_manager.py @@ -0,0 +1,301 @@ +#!/usr/bin/env python3 +""" +Schema Manager for A-Lab Pipeline (Pydantic-First, Auto-Discovery) + +Manages Pydantic schema classes for all parquet tables. +Schemas are defined directly in Python, NOT generated from YAML. + +EXTENSION: To add a new schema (e.g., SEM data): +1. Create a new file: data/products/schema/sem_data.py +2. Define a Pydantic class with __schema_table__ = "sem_data" +3. The schema will be auto-discovered on next pipeline run + +Schema files are located in data/products/schema/: +- experiments.py (main table, consolidated) +- experiment_elements.py +- powder_doses.py +- temperature_logs.py +- xrd_data_points.py +- workflow_tasks.py +- xrd_refinements.py +- xrd_phases.py +- [your_new_schema.py] - auto-discovered! +""" + +import importlib.util +import inspect +import logging +from pathlib import Path +from typing import Dict, List, Optional, Type +from pydantic import BaseModel + +logger = logging.getLogger(__name__) + +SCHEMA_DIR = Path(__file__).parent / "schema" + +# Files to skip during auto-discovery +SKIP_FILES = {'__init__.py', 'base.py', '__pycache__'} + +# Fallback mapping for files without __schema_table__ attribute +# Maps filename (without .py) to expected class name +LEGACY_SCHEMA_MAPPING = { + 'experiments': 'Experiment', + 'experiment_elements': 'ExperimentElement', + 'powder_doses': 'PowderDose', + 'temperature_logs': 'TemperatureLogEntry', + 'workflow_tasks': 'WorkflowTask', + 'xrd_data_points': 'XRDDataPoint', + 'xrd_refinements': 'XRDRefinement', + 'xrd_phases': 'XRDPhase', + 'heating': 'HeatingResult', + 'recovery': 'RecoverPowderResult', + 'diffraction': 'DiffractionResult', + 'powder_dosing': 'PowderDosingSampleResult', +} + + +class SchemaManager: + """ + Manages all Pydantic schema classes for the A-Lab pipeline. + + Auto-discovers schemas from the schema directory. To add a new schema: + 1. Create a .py file in data/products/schema/ + 2. Define a Pydantic BaseModel class + 3. Optionally add __schema_table__ = "table_name" to specify the table name + (otherwise derived from filename) + """ + + def __init__(self, schema_dir: Path = SCHEMA_DIR): + self.schema_dir = Path(schema_dir) + self.schemas: Dict[str, Type[BaseModel]] = {} + self._discover_schemas() + + def _discover_schemas(self): + """ + Auto-discover all Pydantic schema classes from the schema directory. + + Discovery rules: + 1. Scan all .py files in schema_dir (except SKIP_FILES) + 2. Look for classes that inherit from BaseModel + 3. Use __schema_table__ attribute if present, else derive from filename + 4. Fall back to LEGACY_SCHEMA_MAPPING for backwards compatibility + """ + if not self.schema_dir.exists(): + logger.warning(f"Schema directory not found: {self.schema_dir}") + return + + # Auto-discover all .py files + for schema_file in sorted(self.schema_dir.glob("*.py")): + if schema_file.name in SKIP_FILES: + continue + + table_name = schema_file.stem # filename without .py + + try: + # Import the module + spec = importlib.util.spec_from_file_location(table_name, schema_file) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + # Find BaseModel subclasses in the module + schema_class = self._find_schema_class(module, table_name) + + if schema_class: + # Get table name from class attribute or filename + actual_table_name = getattr(schema_class, '__schema_table__', table_name) + self.schemas[actual_table_name] = schema_class + logger.debug(f"Discovered schema: {actual_table_name} -> {schema_class.__name__}") + + except Exception as e: + logger.error(f"Failed to load schema {schema_file}: {e}") + + logger.info(f"Auto-discovered {len(self.schemas)} schemas from {self.schema_dir}") + + def _find_schema_class(self, module, table_name: str) -> Optional[Type[BaseModel]]: + """ + Find the main schema class in a module. + + Priority: + 1. Class with __schema_table__ attribute matching table_name + 2. Class name from LEGACY_SCHEMA_MAPPING + 3. First BaseModel subclass found (excluding imports) + """ + candidates = [] + + for name, obj in inspect.getmembers(module, inspect.isclass): + # Skip imported classes (only want classes defined in this module) + if obj.__module__ != module.__name__: + continue + + # Must be a BaseModel subclass + if not issubclass(obj, BaseModel) or obj is BaseModel: + continue + + # Check for explicit table name + if hasattr(obj, '__schema_table__') and obj.__schema_table__ == table_name: + return obj + + candidates.append((name, obj)) + + # Try legacy mapping + if table_name in LEGACY_SCHEMA_MAPPING: + expected_class = LEGACY_SCHEMA_MAPPING[table_name] + for name, obj in candidates: + if name == expected_class: + return obj + + # Return first candidate if any + if candidates: + return candidates[0][1] + + return None + + def get_schema(self, table_name: str) -> Optional[Type[BaseModel]]: + """Get Pydantic schema class for a specific table""" + return self.schemas.get(table_name) + + def get_all_schemas(self) -> Dict[str, Type[BaseModel]]: + """Get all loaded schema classes""" + return self.schemas + + def get_table_names(self) -> List[str]: + """Get list of all table names""" + return sorted(self.schemas.keys()) + + def get_main_schema(self) -> Optional[Type[BaseModel]]: + """Get the main Experiment schema (consolidated table)""" + return self.schemas.get('experiments') + + def get_uploadable_fields(self, table_name: str) -> List[str]: + """ + Get list of fields that should be uploaded to MPContribs. + + Fields marked with exclude_from_upload=True are excluded. + """ + schema = self.get_schema(table_name) + if not schema: + return [] + + uploadable = [] + for field_name, field_info in schema.model_fields.items(): + extra = field_info.json_schema_extra or {} + if not extra.get("exclude_from_upload", False): + uploadable.append(field_name) + + return uploadable + + def get_excluded_fields(self, table_name: str) -> List[str]: + """ + Get list of fields that should NOT be uploaded to MPContribs. + """ + schema = self.get_schema(table_name) + if not schema: + return [] + + excluded = [] + for field_name, field_info in schema.model_fields.items(): + extra = field_info.json_schema_extra or {} + if extra.get("exclude_from_upload", False): + excluded.append(field_name) + + return excluded + + def validate_row(self, table_name: str, row_data: dict) -> tuple[bool, Optional[str]]: + """ + Validate a single row of data against its Pydantic schema. + + Args: + table_name: Name of the table + row_data: Dictionary of field values + + Returns: + (is_valid, error_message) + """ + schema = self.get_schema(table_name) + if not schema: + return True, None # No schema = skip validation + + try: + schema(**row_data) + return True, None + except Exception as e: + return False, str(e) + + def get_schema_fields(self, table_name: str) -> Dict[str, dict]: + """ + Get field information for a schema. + + Returns dict of field_name -> {type, required, description} + """ + schema = self.get_schema(table_name) + if not schema: + return {} + + fields = {} + for field_name, field_info in schema.model_fields.items(): + fields[field_name] = { + 'type': str(field_info.annotation), + 'required': field_info.is_required(), + 'description': field_info.description or '', + 'exclude_from_upload': (field_info.json_schema_extra or {}).get('exclude_from_upload', False) + } + + return fields + + def get_schema_summary(self) -> Dict[str, Dict]: + """Get summary of all schemas""" + summary = {} + for table_name, schema in self.schemas.items(): + doc = schema.__doc__ or '' + summary[table_name] = { + 'description': doc.split('\n')[0].strip() if doc else '', + 'class_name': schema.__name__, + 'num_fields': len(schema.model_fields), + 'uploadable_fields': len(self.get_uploadable_fields(table_name)), + 'excluded_fields': len(self.get_excluded_fields(table_name)) + } + return summary + + +if __name__ == '__main__': + # Test the schema manager with auto-discovery + from rich.console import Console + from rich.table import Table + + console = Console() + + manager = SchemaManager() + + console.print(f"\n[bold cyan]Schema Auto-Discovery Results[/bold cyan]") + console.print(f"Directory: {manager.schema_dir}") + console.print(f"Discovered: {len(manager.schemas)} schemas\n") + + # Create table + table = Table(title="Discovered Schemas") + table.add_column("Table Name", style="cyan") + table.add_column("Class", style="green") + table.add_column("Fields", justify="right") + table.add_column("Uploadable", justify="right") + table.add_column("Excluded", style="yellow") + + for table_name in manager.get_table_names(): + schema = manager.get_schema(table_name) + excluded = manager.get_excluded_fields(table_name) + uploadable = manager.get_uploadable_fields(table_name) + + table.add_row( + table_name, + schema.__name__, + str(len(schema.model_fields)), + str(len(uploadable)), + ", ".join(excluded) if excluded else "-" + ) + + console.print(table) + + console.print("\n[bold green]✓ Auto-discovery complete![/bold green]") + console.print("\n[dim]To add a new schema:[/dim]") + console.print(" 1. Create: data/products/schema/your_schema.py") + console.print(" 2. Define: class YourSchema(BaseModel, extra='forbid'): ...") + console.print(" 3. Optional: __schema_table__ = 'your_table_name'") + console.print(" 4. Run pipeline - schema auto-discovered!") diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/schema_validator.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/schema_validator.py new file mode 100644 index 000000000..cbe0deabe --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/schema_validator.py @@ -0,0 +1,217 @@ +#!/usr/bin/env python3 +""" +Schema Validator - Validates parquet data against Pydantic schemas + +Uses the Pydantic schema classes directly for validation. +No YAML files - schemas are Python-first. +""" + +import logging +from pathlib import Path +from typing import Dict, List, Tuple, Type +import pandas as pd +from pydantic import BaseModel + +logger = logging.getLogger(__name__) + + +class SchemaValidator: + """Validates parquet data against Pydantic schemas""" + + def __init__(self, schema_manager=None): + """ + Args: + schema_manager: SchemaManager instance (will create if not provided) + """ + if schema_manager is None: + from schema_manager import SchemaManager + self.schema_manager = SchemaManager() + else: + self.schema_manager = schema_manager + + self.warnings = [] + self.errors = [] + self.validation_stats = {} + + def validate_parquet_data(self, parquet_dir: Path) -> Tuple[bool, List[str], List[str]]: + """ + Validate all parquet files against their Pydantic schemas. + + Args: + parquet_dir: Directory containing parquet files + + Returns: + (is_valid, errors, warnings) + """ + self.warnings = [] + self.errors = [] + self.validation_stats = {} + + parquet_dir = Path(parquet_dir) + + # Validate each table + for table_name in self.schema_manager.get_table_names(): + schema = self.schema_manager.get_schema(table_name) + + if not schema: + continue + + parquet_file = parquet_dir / f"{table_name}.parquet" + + if not parquet_file.exists(): + # Temperature logs and XRD data points are optional + if table_name in ['temperature_logs', 'xrd_data_points']: + logger.debug(f"Optional table not present: {table_name}.parquet") + else: + self.warnings.append(f"Table not found: {table_name}.parquet") + continue + + # Validate this table + self._validate_table(table_name, parquet_file, schema) + + is_valid = len(self.errors) == 0 + return is_valid, self.errors, self.warnings + + def _validate_table(self, table_name: str, parquet_file: Path, schema: Type[BaseModel]): + """ + Validate a parquet table against its Pydantic schema. + + Args: + table_name: Name of the table + parquet_file: Path to the parquet file + schema: Pydantic schema class + """ + try: + df = pd.read_parquet(parquet_file) + except Exception as e: + self.errors.append(f"[{table_name}] Failed to load parquet: {e}") + return + + # Check column existence + schema_fields = set(schema.model_fields.keys()) + parquet_columns = set(df.columns) + + missing_required = [] + for field_name, field_info in schema.model_fields.items(): + if field_info.is_required() and field_name not in parquet_columns: + missing_required.append(field_name) + + if missing_required: + self.errors.append( + f"[{table_name}] Missing required columns: {', '.join(missing_required)}" + ) + + # Missing optional columns are just warnings + missing_optional = schema_fields - parquet_columns - set(missing_required) + if missing_optional: + self.warnings.append( + f"[{table_name}] Missing optional columns: {', '.join(sorted(missing_optional)[:5])}" + + (f" (+{len(missing_optional)-5} more)" if len(missing_optional) > 5 else "") + ) + + # Validate a sample of rows with Pydantic + sample_size = min(100, len(df)) + sample_df = df.sample(n=sample_size, random_state=42) if len(df) > sample_size else df + + valid_count = 0 + error_count = 0 + error_samples = [] + + for idx, row in sample_df.iterrows(): + try: + row_dict = {k: v for k, v in row.to_dict().items() if pd.notna(v) or k in schema_fields} + # Handle NaN/None values + cleaned_dict = {} + for k, v in row_dict.items(): + if pd.isna(v): + cleaned_dict[k] = None + else: + cleaned_dict[k] = v + + schema(**cleaned_dict) + valid_count += 1 + except Exception as e: + error_count += 1 + if len(error_samples) < 3: + error_samples.append(f"Row {idx}: {str(e)[:150]}") + + # Store stats + self.validation_stats[table_name] = { + 'total': len(df), + 'sampled': sample_size, + 'valid': valid_count, + 'errors': error_count + } + + # Report validation errors + if error_count > 0: + self.errors.append( + f"[{table_name}] {error_count}/{sample_size} sampled rows failed validation" + ) + for sample in error_samples: + self.errors.append(f" • {sample}") + else: + logger.debug(f"[{table_name}] ✓ {valid_count} sampled rows valid") + + +def validate_product_schema(product_dir: Path, schema_manager=None) -> bool: + """ + Validate a product's parquet data against Pydantic schemas. + + Args: + product_dir: Product directory containing parquet_data/ + schema_manager: Optional SchemaManager instance + + Returns: + True if validation passes (or passes with warnings only) + """ + from rich.console import Console + console = Console() + + validator = SchemaValidator(schema_manager) + parquet_dir = product_dir / "parquet_data" + + if not parquet_dir.exists(): + console.print(f"[yellow]⚠ No parquet data found at {parquet_dir}[/yellow]") + return True + + is_valid, errors, warnings = validator.validate_parquet_data(parquet_dir) + + # Report validation statistics + if validator.validation_stats: + console.print("\n[bold]Validation Statistics:[/bold]") + for table_name, stats in validator.validation_stats.items(): + if stats['errors'] == 0: + console.print(f" [green]✓ {table_name}: {stats['valid']}/{stats['sampled']} sampled valid ({stats['total']} total)[/green]") + else: + console.print(f" [red]✗ {table_name}: {stats['valid']}/{stats['sampled']} sampled valid ({stats['errors']} errors)[/red]") + + if errors: + console.print("\n[red]✗ Schema Validation Errors:[/red]") + for error in errors: + console.print(f" [red]• {error}[/red]") + + if warnings: + console.print("\n[yellow]⚠ Schema Validation Warnings:[/yellow]") + for warning in warnings: + console.print(f" [yellow]• {warning}[/yellow]") + + if is_valid and not warnings: + console.print("\n[green]✓ Schema validation passed![/green]") + elif is_valid: + console.print("\n[green]✓ Schema validation passed with warnings[/green]") + + return is_valid + + +if __name__ == "__main__": + import sys + if len(sys.argv) < 2: + print("Usage: python schema_validator.py ") + sys.exit(1) + + product_name = sys.argv[1] + product_dir = Path(f"data/products/{product_name}") + + success = validate_product_schema(product_dir) + sys.exit(0 if success else 1) diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/run_product_pipeline.sh b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/run_product_pipeline.sh new file mode 100755 index 000000000..f6c15b72b --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/run_product_pipeline.sh @@ -0,0 +1,38 @@ +#!/bin/bash +# +# Product Pipeline Runner - Uses existing data/venv +# +# Usage: +# ./run_product_pipeline.sh create # Create new product +# ./run_product_pipeline.sh list # List products +# ./run_product_pipeline.sh run --product # Run pipeline (dry run) +# ./run_product_pipeline.sh run --product --upload # Upload (MPContribs + S3) +# ./run_product_pipeline.sh status --product # Check status +# + +set -e + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +# Check if venv exists +if [ ! -d "data/venv" ]; then + echo "→ Creating virtual environment..." + python3 -m venv data/venv + echo "✓ Virtual environment created" +fi + +# Activate venv +source data/venv/bin/activate + +# Install/update dependencies +echo "→ Checking dependencies..." +pip install -q -r data/requirements.txt 2>&1 | grep -v "UserWarning" | grep -v "Valid config keys" || true + +# Run the pipeline CLI (default to 'create' if no args) +if [ $# -eq 0 ]; then + python data/pipeline/run_pipeline.py create +else +python data/pipeline/run_pipeline.py "$@" +fi + From 67eda5a7b0cefdb1fbd73f44239eb5b75692c976 Mon Sep 17 00:00:00 2001 From: Silen Naihin Date: Thu, 18 Dec 2025 15:49:13 -0500 Subject: [PATCH 3/3] link to repo as source of truth --- .../lux/projects/alab/pipelines/README.md | 13 + .../pipelines/data/analyses/base_analyzer.py | 479 --------- .../alab/pipelines/data/config/README.md | 211 ---- .../alab/pipelines/data/config/analyses.yaml | 135 --- .../pipelines/data/config/config_loader.py | 327 ------- .../alab/pipelines/data/config/defaults.yaml | 94 -- .../alab/pipelines/data/config/filters.yaml | 97 -- .../alab/pipelines/data/mongodb_to_parquet.py | 591 ----------- .../data/pipeline/PIPELINE_ARCHITECTURE.md | 472 --------- .../alab/pipelines/data/pipeline/README.md | 267 ----- .../alab/pipelines/data/pipeline/__init__.py | 18 - .../data/pipeline/mpcontribs_setup.py | 334 ------- .../pipelines/data/pipeline/pipeline_state.py | 292 ------ .../data/pipeline/product_pipeline.py | 691 ------------- .../pipelines/data/pipeline/run_pipeline.py | 429 -------- .../pipelines/data/pipeline/s3_uploader.py | 420 -------- .../data/pipeline/test_integrated_pipeline.py | 920 ------------------ .../pipelines/data/products/SCHEMA_SYSTEM.md | 202 ---- .../pipelines/data/products/base_product.py | 771 --------------- .../pipelines/data/products/schema_manager.py | 301 ------ .../data/products/schema_validator.py | 217 ----- .../alab/pipelines/run_product_pipeline.sh | 38 - 22 files changed, 13 insertions(+), 7306 deletions(-) create mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/README.md delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/analyses/base_analyzer.py delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/README.md delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/analyses.yaml delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/config_loader.py delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/defaults.yaml delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/filters.yaml delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/mongodb_to_parquet.py delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/PIPELINE_ARCHITECTURE.md delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/README.md delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/__init__.py delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/mpcontribs_setup.py delete mode 100755 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/pipeline_state.py delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/product_pipeline.py delete mode 100755 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/run_pipeline.py delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/s3_uploader.py delete mode 100755 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/test_integrated_pipeline.py delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/SCHEMA_SYSTEM.md delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/base_product.py delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/schema_manager.py delete mode 100644 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/schema_validator.py delete mode 100755 mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/run_product_pipeline.sh diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/README.md b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/README.md new file mode 100644 index 000000000..fda81bfe2 --- /dev/null +++ b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/README.md @@ -0,0 +1,13 @@ +# A-Lab Pipelines + +Pipeline code for the A-Lab project is maintained in a separate repository: + +**Repository:** [alab-pipeline](https://github.com/SilenNaihin/alab-pipeline.git) + +## Getting Started + +To clone the pipeline repository: + +```bash +git clone https://github.com/SilenNaihin/alab-pipeline.git +``` diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/analyses/base_analyzer.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/analyses/base_analyzer.py deleted file mode 100644 index 72d92468b..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/analyses/base_analyzer.py +++ /dev/null @@ -1,479 +0,0 @@ -#!/usr/bin/env python3 -""" -Base Analysis Plugin Interface (Auto-Discovery) - -All analysis modules should inherit from BaseAnalyzer and implement: -1. analyze() - Run the analysis on experiments -2. get_output_schema() - Define output schema for validation - -EXTENSION: To add a new analysis: -1. Create a new file: data/analyses/your_analysis.py -2. Define a class inheriting from BaseAnalyzer -3. Set class attributes: - - name: str = "your_analysis" (used in product config) - - description: str = "What it does" - - cli_flag: str = "--your-analysis" (optional, for CLI) -4. The analysis will be auto-discovered on next pipeline run - -Analysis modules are auto-discovered from data/analyses/*.py -""" - -from abc import ABC, abstractmethod -from pathlib import Path -from typing import Dict, List, Optional, Any, Type -import pandas as pd -import logging -import importlib.util -import inspect - -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) - -# Directory containing analysis plugins -ANALYSES_DIR = Path(__file__).parent - -# Files to skip during auto-discovery -SKIP_FILES = {'__init__.py', 'base_analyzer.py', '__pycache__'} - - -class BaseAnalyzer(ABC): - """ - Base class for all analysis modules. - - To create a new analyzer: - 1. Inherit from BaseAnalyzer - 2. Set class attributes: name, description (optional: cli_flag) - 3. Implement analyze() and get_output_schema() - - Example: - class SEMAnalyzer(BaseAnalyzer): - name = "sem_clustering" - description = "Cluster SEM images by morphology" - cli_flag = "--sem" - - def analyze(self, experiments_df, parquet_dir): - # Your analysis logic - return results_df - - def get_output_schema(self): - return {'cluster_id': {'type': 'int', 'required': True}} - """ - - # Class attributes for discovery (override in subclasses) - name: str = None # Required: unique identifier for this analyzer - description: str = "" # Optional: human-readable description - cli_flag: str = None # Optional: command line flag (e.g., "--xrd") - - def __init__(self, config: Dict = None): - """ - Initialize analyzer with configuration - - Args: - config: Analysis-specific configuration from product config - """ - self.config = config or {} - # Use class name attribute or fall back to class name - if self.name is None: - self.name = self.__class__.__name__.lower().replace('analyzer', '') - - @abstractmethod - def analyze(self, - experiments_df: pd.DataFrame, - parquet_dir: Path) -> pd.DataFrame: - """ - Run analysis on experiments - - Args: - experiments_df: DataFrame of experiments to analyze - parquet_dir: Directory containing parquet files - - Returns: - DataFrame with analysis results (one row per experiment) - """ - pass - - @abstractmethod - def get_output_schema(self) -> Dict[str, Dict]: - """ - Get output schema for validation - - Returns: - Dict of field_name -> {type, description, required} - """ - pass - - def validate_output(self, results_df: pd.DataFrame) -> bool: - """ - Validate that results match expected schema - - Args: - results_df: Analysis results DataFrame - - Returns: - True if valid - """ - schema = self.get_output_schema() - - for field, spec in schema.items(): - if spec.get('required', False) and field not in results_df.columns: - logger.error(f"Missing required field: {field}") - return False - - if field in results_df.columns: - # Type checking could be added here - pass - - return True - - def save_results(self, results_df: pd.DataFrame, output_dir: Path): - """Save analysis results to parquet""" - output_dir.mkdir(parents=True, exist_ok=True) - output_file = output_dir / f"{self.name.lower()}_results.parquet" - results_df.to_parquet(output_file, index=False) - logger.info(f"Saved {len(results_df)} results to {output_file}") - - -class XRDAnalyzer(BaseAnalyzer): - """XRD Phase Analysis using DARA""" - - # Class attributes for discovery - name = "xrd_dara" - description = "XRD phase identification using DARA (Deep Analysis for Rietveld Automation)" - cli_flag = "--xrd" - - def __init__(self, config: Dict = None): - super().__init__(config) - - def analyze(self, experiments_df: pd.DataFrame, parquet_dir: Path) -> pd.DataFrame: - """Run DARA XRD analysis on experiments""" - import subprocess - import json - - results = [] - xrd_results_dir = Path("data/xrd_creation/results") - - for _, exp in experiments_df.iterrows(): - exp_name = exp['name'] - result_file = xrd_results_dir / f"{exp_name}_result.json" - - # Check if already analyzed - if result_file.exists(): - with open(result_file, 'r') as f: - result = json.load(f) - results.append({ - 'experiment_name': exp_name, - 'xrd_success': result.get('success', False), - 'xrd_rwp': result.get('rwp'), - 'xrd_num_phases': result.get('num_phases', 0), - 'xrd_error': result.get('error') - }) - else: - # Run analysis - try: - cmd = [ - "python", - "data/xrd_creation/analyze_single.py", - "--experiment", exp_name, - "--mode", "phase_search" - ] - - # Add config options - if self.config.get('wmin'): - cmd.extend(["--wmin", str(self.config['wmin'])]) - if self.config.get('wmax'): - cmd.extend(["--wmax", str(self.config['wmax'])]) - - result = subprocess.run(cmd, capture_output=True, text=True, timeout=120) - - # Load result - if result_file.exists(): - with open(result_file, 'r') as f: - analysis_result = json.load(f) - results.append({ - 'experiment_name': exp_name, - 'xrd_success': analysis_result.get('success', False), - 'xrd_rwp': analysis_result.get('rwp'), - 'xrd_num_phases': analysis_result.get('num_phases', 0), - 'xrd_error': analysis_result.get('error') - }) - else: - results.append({ - 'experiment_name': exp_name, - 'xrd_success': False, - 'xrd_rwp': None, - 'xrd_num_phases': 0, - 'xrd_error': 'Analysis failed' - }) - - except Exception as e: - logger.error(f"XRD analysis failed for {exp_name}: {e}") - results.append({ - 'experiment_name': exp_name, - 'xrd_success': False, - 'xrd_rwp': None, - 'xrd_num_phases': 0, - 'xrd_error': str(e) - }) - - return pd.DataFrame(results) - - def get_output_schema(self) -> Dict[str, Dict]: - return { - 'xrd_success': {'type': 'boolean', 'required': True, 'description': 'XRD analysis succeeded'}, - 'xrd_rwp': {'type': 'float', 'required': False, 'description': 'Weighted profile R-factor'}, - 'xrd_num_phases': {'type': 'int', 'required': True, 'description': 'Number of phases identified'}, - 'xrd_error': {'type': 'string', 'required': False, 'description': 'Error message if failed'} - } - - -class PowderStatisticsAnalyzer(BaseAnalyzer): - """Analyze powder dosing statistics""" - - # Class attributes for discovery - name = "powder_statistics" - description = "Calculate powder dosing accuracy and statistics" - cli_flag = "--powder-stats" - - def __init__(self, config: Dict = None): - super().__init__(config) - - def analyze(self, experiments_df: pd.DataFrame, parquet_dir: Path) -> pd.DataFrame: - """Calculate powder dosing statistics""" - - # Load powder dosing data - powder_doses_df = pd.read_parquet(parquet_dir / "powder_doses.parquet") - - results = [] - - for _, exp in experiments_df.iterrows(): - exp_id = exp['experiment_id'] - exp_name = exp['name'] - - # Get doses for this experiment - exp_doses = powder_doses_df[powder_doses_df['experiment_id'] == exp_id] - - if len(exp_doses) > 0: - # Calculate statistics - avg_accuracy = exp_doses['accuracy_percent'].mean() if 'accuracy_percent' in exp_doses.columns else None - total_doses = len(exp_doses) - unique_powders = exp_doses['powder_name'].nunique() - total_mass = exp_doses['actual_mass'].sum() - - results.append({ - 'experiment_name': exp_name, - 'powder_avg_accuracy': avg_accuracy, - 'powder_total_doses': total_doses, - 'powder_unique_count': unique_powders, - 'powder_total_mass_g': total_mass - }) - else: - results.append({ - 'experiment_name': exp_name, - 'powder_avg_accuracy': None, - 'powder_total_doses': 0, - 'powder_unique_count': 0, - 'powder_total_mass_g': 0.0 - }) - - return pd.DataFrame(results) - - def get_output_schema(self) -> Dict[str, Dict]: - return { - 'powder_avg_accuracy': {'type': 'float', 'required': False, 'description': 'Average dosing accuracy %'}, - 'powder_total_doses': {'type': 'int', 'required': True, 'description': 'Total number of doses'}, - 'powder_unique_count': {'type': 'int', 'required': True, 'description': 'Number of unique powders'}, - 'powder_total_mass_g': {'type': 'float', 'required': True, 'description': 'Total powder mass in grams'} - } - - -class AnalysisPluginManager: - """ - Manages and auto-discovers analysis plugins. - - Analyzers are discovered from: - 1. Built-in analyzers in this file (XRDAnalyzer, PowderStatisticsAnalyzer) - 2. Any .py file in data/analyses/ containing a BaseAnalyzer subclass - - To add a new analyzer: - 1. Create data/analyses/your_analyzer.py - 2. Define class YourAnalyzer(BaseAnalyzer) with name attribute - 3. It will be auto-discovered - """ - - def __init__(self, analyses_dir: Path = None): - self.analyses_dir = Path(analyses_dir or ANALYSES_DIR) - self.analyzers: Dict[str, Type[BaseAnalyzer]] = {} - self._discover_analyzers() - - def _discover_analyzers(self): - """ - Auto-discover available analyzers from the analyses directory. - - Discovery rules: - 1. Register built-in analyzers first - 2. Scan all .py files in analyses_dir (except SKIP_FILES) - 3. Find classes that inherit from BaseAnalyzer - 4. Register using the class's 'name' attribute - """ - # Built-in analyzers (defined in this file) - self.analyzers['xrd_dara'] = XRDAnalyzer - self.analyzers['powder_statistics'] = PowderStatisticsAnalyzer - - # Auto-discover from analyses directory - if not self.analyses_dir.exists(): - logger.warning(f"Analyses directory not found: {self.analyses_dir}") - return - - for analyzer_file in sorted(self.analyses_dir.glob("*.py")): - if analyzer_file.name in SKIP_FILES: - continue - - try: - # Import the module - module_name = analyzer_file.stem - spec = importlib.util.spec_from_file_location(module_name, analyzer_file) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - - # Find BaseAnalyzer subclasses - for name, obj in inspect.getmembers(module, inspect.isclass): - # Skip imported classes and BaseAnalyzer itself - if obj.__module__ != module.__name__: - continue - if not issubclass(obj, BaseAnalyzer) or obj is BaseAnalyzer: - continue - - # Get analyzer name from class attribute - analyzer_name = getattr(obj, 'name', None) - if analyzer_name is None: - # Derive from class name: SEMAnalyzer -> sem - analyzer_name = name.lower().replace('analyzer', '') - - # Don't override built-in analyzers - if analyzer_name not in self.analyzers: - self.analyzers[analyzer_name] = obj - logger.debug(f"Discovered analyzer: {analyzer_name} -> {name}") - - except Exception as e: - logger.error(f"Failed to load analyzer from {analyzer_file}: {e}") - - logger.info(f"Discovered {len(self.analyzers)} analyzers: {', '.join(self.analyzers.keys())}") - - def get_analyzer(self, name: str, config: Dict = None) -> Optional[BaseAnalyzer]: - """Get analyzer instance by name""" - if name in self.analyzers: - return self.analyzers[name](config) - - logger.warning(f"Analyzer '{name}' not found. Available: {', '.join(self.analyzers.keys())}") - return None - - def list_analyzers(self) -> List[str]: - """List available analyzer names""" - return sorted(self.analyzers.keys()) - - def get_analyzer_info(self) -> Dict[str, Dict]: - """Get info about all available analyzers""" - info = {} - for name, analyzer_class in self.analyzers.items(): - info[name] = { - 'class': analyzer_class.__name__, - 'description': getattr(analyzer_class, 'description', ''), - 'cli_flag': getattr(analyzer_class, 'cli_flag', None), - } - return info - - def run_analyses(self, - analyses: List[Dict], - experiments_df: pd.DataFrame, - parquet_dir: Path, - output_dir: Path) -> pd.DataFrame: - """ - Run multiple analyses and merge results - - Args: - analyses: List of {name, config} dicts - experiments_df: Experiments to analyze - parquet_dir: Input parquet directory - output_dir: Output directory for results - - Returns: - Merged DataFrame with all analysis results - """ - # Start with essential experiment fields - base_fields = ['experiment_id', 'name', 'experiment_type', 'target_formula'] - # Only include fields that exist in the DataFrame - available_fields = [f for f in base_fields if f in experiments_df.columns] - all_results = experiments_df[available_fields].copy() - - for analysis_config in analyses: - if not analysis_config.get('enabled', True): - continue - - analyzer = self.get_analyzer( - analysis_config['name'], - analysis_config.get('config', {}) - ) - - if analyzer: - logger.info(f"Running {analyzer.name} analysis...") - - try: - results = analyzer.analyze(experiments_df, parquet_dir) - - if analyzer.validate_output(results): - # Merge results - all_results = all_results.merge( - results, - left_on='name', - right_on='experiment_name', - how='left' - ) - - # Save individual results - analyzer.save_results(results, output_dir) - else: - logger.error(f"Validation failed for {analyzer.name}") - - except Exception as e: - logger.error(f"Analysis {analyzer.name} failed: {e}") - - return all_results - - -# CLI for listing available analyzers -if __name__ == '__main__': - from rich.console import Console - from rich.table import Table - - console = Console() - - manager = AnalysisPluginManager() - - console.print(f"\n[bold cyan]Analysis Plugin Auto-Discovery[/bold cyan]") - console.print(f"Directory: {manager.analyses_dir}") - console.print(f"Discovered: {len(manager.analyzers)} analyzers\n") - - # Create table - table = Table(title="Available Analyzers") - table.add_column("Name", style="cyan") - table.add_column("Class", style="green") - table.add_column("CLI Flag", style="yellow") - table.add_column("Description") - - for name, info in manager.get_analyzer_info().items(): - table.add_row( - name, - info['class'], - info['cli_flag'] or "-", - info['description'] or "-" - ) - - console.print(table) - - console.print("\n[bold green]✓ Auto-discovery complete![/bold green]") - console.print("\n[dim]To add a new analyzer:[/dim]") - console.print(" 1. Create: data/analyses/your_analyzer.py") - console.print(" 2. Define: class YourAnalyzer(BaseAnalyzer): ...") - console.print(" 3. Set: name = 'your_analyzer'") - console.print(" 4. Implement: analyze() and get_output_schema()") - console.print(" 5. Run pipeline - analyzer auto-discovered!") diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/README.md b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/README.md deleted file mode 100644 index e7a8b7005..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/README.md +++ /dev/null @@ -1,211 +0,0 @@ -# A-Lab Pipeline Configuration - -Configuration system with three-layer priority: - -``` -1. Environment Variables (highest) → 2. YAML Files → 3. Code Defaults (fallback) -``` - -## Quick Start - -### Using Environment Variables (Recommended for Production) - -**Option 1: Copy example file** - -```bash -# Copy the example .env file (with current defaults) -cp data/config/env.example .env - -# Edit .env and uncomment/modify values -# Then source it before running -source .env -./update_data.sh -``` - -**Option 2: Export directly** - -```bash -# Set MongoDB connection -export ALAB_MONGO_URI="mongodb://production-host:27017/" -export ALAB_MONGO_DB="production" - -# Set S3 bucket -export ALAB_S3_BUCKET="my-custom-bucket" - -# Run pipeline (uses env vars automatically) -./update_data.sh -``` - -### Using YAML Files (Recommended for Development) - -Edit `data/config/defaults.yaml`: - -```yaml -mongodb: - uri: 'mongodb://localhost:27017/' - database: 'my_database' - collection: 'my_collection' -``` - -### View Current Configuration - -```bash -python data/config/config_loader.py -``` - -Shows all loaded values and their sources (env vs yaml vs defaults). - -## Configuration Files - -| File | Purpose | -| -------------------- | ------------------------------- | -| **defaults.yaml** | Global pipeline defaults | -| **filters.yaml** | Experiment filter presets | -| **analyses.yaml** | Analysis plugin documentation | -| **config_loader.py** | Configuration loading system | -| **env.example** | Example env file (copy to .env) | - -## Environment Variables - -### MongoDB - -```bash -ALAB_MONGO_URI=mongodb://localhost:27017/ # MongoDB connection URI -ALAB_MONGO_DB=temporary # Database name -ALAB_MONGO_COLLECTION=release # Collection name -``` - -### S3 Upload - -```bash -ALAB_S3_BUCKET=materialsproject-contribs # S3 bucket name -ALAB_S3_PREFIX=alab_synthesis # S3 prefix path -ALAB_S3_EXCLUDE_LARGE=true # Exclude large files -ALAB_S3_LARGE_THRESHOLD_MB=50 # Large file threshold (MB) -``` - -### Parquet Options - -```bash -ALAB_SKIP_TEMP_LOGS=false # Skip temperature logs -ALAB_SKIP_XRD_POINTS=false # Skip XRD data points -ALAB_SKIP_WORKFLOW_TASKS=false # Skip workflow tasks -ALAB_PARQUET_COMPRESSION=snappy # Compression: snappy, gzip, brotli -ALAB_PARQUET_ENGINE=pyarrow # Engine: pyarrow, fastparquet -``` - -### Materials Project API - -```bash -ALAB_MP_API_KEY=your_api_key # MP API key (for XRD analysis) -# OR -MP_API_KEY=your_api_key # Alternative name -``` - -## Usage in Scripts - -### Python - -```python -from config_loader import get_config - -# Get configuration -config = get_config() - -# Access values -print(config.mongo_uri) # mongodb://localhost:27017/ -print(config.mongo_db) # temporary -print(config.s3_bucket) # materialsproject-contribs - -# Or use convenience functions -from config_loader import get_mongo_uri, get_s3_bucket - -uri = get_mongo_uri() # Gets from env > yaml > default -bucket = get_s3_bucket() -``` - -### Shell Scripts - -```bash -# Use environment variables directly -: ${ALAB_MONGO_URI:="mongodb://localhost:27017/"} - -# Or source from .env file -if [ -f data/.env ]; then - export $(grep -v '^#' data/.env | xargs) -fi -``` - -## Configuration Priority Examples - -### Example 1: All from YAML - -```bash -# No env vars set -$ python data/config/config_loader.py -MongoDB URI: mongodb://localhost:27017/ (from YAML) -``` - -### Example 2: Override with Env - -```bash -# Set env var -$ export ALAB_MONGO_URI="mongodb://production:27017/" -$ python data/config/config_loader.py -MongoDB URI: mongodb://production:27017/ (from ENV) ✓ -``` - -### Example 3: Mixed Sources - -```bash -# Some from env, some from yaml -$ export ALAB_MONGO_URI="mongodb://prod:27017/" # Custom URI -# Leave ALAB_MONGO_DB unset # Use YAML default -$ python data/config/config_loader.py -MongoDB URI: mongodb://prod:27017/ (from ENV) ✓ -MongoDB DB: temporary (from YAML) -``` - -## Best Practices - -1. **Development**: Use `defaults.yaml` for local development -2. **Production**: Use environment variables for sensitive values -3. **Testing**: Use env vars to point to test databases -4. **CI/CD**: Set env vars in your deployment pipeline -5. **Never commit** `.env` files (already in `.gitignore`) - -## Troubleshooting - -### Config not loading? - -```bash -# Check current config -python data/config/config_loader.py - -# Verify env vars are set -env | grep ALAB_ -``` - -### Want to use .env file? - -```bash -# Create from example -cp data/config/env.example .env - -# Edit .env with your values (uncomment lines to override defaults) -nano .env - -# Source it before running scripts -source .env -./update_data.sh -``` - -### Reset to defaults - -```bash -# Unset all ALAB env vars -unset $(env | grep ALAB_ | cut -d= -f1) - -# Now uses YAML/defaults only -./update_data.sh -``` diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/analyses.yaml b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/analyses.yaml deleted file mode 100644 index 5290a7a3f..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/analyses.yaml +++ /dev/null @@ -1,135 +0,0 @@ -# ============================================================================= -# A-Lab Analysis Registry -# ============================================================================= -# Documentation for available analysis plugins. -# Analyses are auto-discovered from data/analyses/*.py -# -# This file serves as: -# 1. Documentation of available analyses -# 2. Default configuration for each analysis -# 3. Template for adding new analyses -# ============================================================================= - -# Built-in analyses (always available) -analyses: - xrd_dara: - description: 'XRD phase identification using DARA' - class: XRDAnalyzer - file: base_analyzer.py - cli_flag: '--xrd' - output_parquet: xrd_refinements.parquet, xrd_phases.parquet - default_config: - wmin: 10 # Minimum 2-theta angle - wmax: 80 # Maximum 2-theta angle - save_viz: false # Save visualization images - outputs: - - xrd_success: 'Whether analysis succeeded' - - xrd_rwp: 'Weighted profile R-factor' - - xrd_num_phases: 'Number of phases identified' - - xrd_error: 'Error message if failed' - requirements: - - experiments.parquet (with xrd_sampleid_in_aeris) - - xrd_data_points.parquet (optional, for patterns) - notes: | - Uses DARA (Deep Analysis for Rietveld Automation) for automated - phase identification. Requires MP API key for CIF downloads. - Results stored in data/xrd_creation/results/ - - powder_statistics: - description: 'Calculate powder dosing statistics' - class: PowderStatisticsAnalyzer - file: base_analyzer.py - cli_flag: '--powder-stats' - output_parquet: null # Results merged into main output - default_config: {} - outputs: - - powder_avg_accuracy: 'Average dosing accuracy %' - - powder_total_doses: 'Total number of doses' - - powder_unique_count: 'Number of unique powders' - - powder_total_mass_g: 'Total powder mass in grams' - requirements: - - experiments.parquet - - powder_doses.parquet - notes: | - Calculates statistics about powder dosing accuracy. - Fast analysis, recommended for all products. - -# ============================================================================= -# Adding a New Analysis -# ============================================================================= -# To add a new analysis (e.g., SEM clustering): -# -# 1. Create the analyzer file: -# data/analyses/sem_analyzer.py -# -# 2. Define the analyzer class: -# ```python -# from base_analyzer import BaseAnalyzer -# -# class SEMAnalyzer(BaseAnalyzer): -# name = "sem_clustering" -# description = "Cluster SEM images by morphology" -# cli_flag = "--sem" -# -# def analyze(self, experiments_df, parquet_dir): -# # Your analysis logic here -# results = [] -# for _, exp in experiments_df.iterrows(): -# # Process each experiment -# results.append({ -# 'experiment_name': exp['name'], -# 'cluster_id': compute_cluster(exp), -# 'morphology_score': compute_score(exp) -# }) -# return pd.DataFrame(results) -# -# def get_output_schema(self): -# return { -# 'cluster_id': {'type': 'int', 'required': True}, -# 'morphology_score': {'type': 'float', 'required': False} -# } -# ``` -# -# 3. Document here (optional): -# sem_clustering: -# description: "Cluster SEM images by morphology" -# class: SEMAnalyzer -# file: sem_analyzer.py -# ... -# -# 4. The analysis will be auto-discovered on next pipeline run -# ============================================================================= - -# Placeholder for future analyses -# Uncomment and modify when adding: - -# sem_clustering: -# description: "Cluster SEM images by morphology" -# class: SEMAnalyzer -# file: sem_analyzer.py -# cli_flag: "--sem" -# default_config: -# num_clusters: 5 -# feature_extraction: "resnet50" -# outputs: -# - cluster_id: "Cluster assignment" -# - morphology_score: "Morphology similarity score" -# requirements: -# - SEM images in experiments/*/SEM images/ -# notes: | -# Uses computer vision to cluster SEM images. -# Requires tensorflow/pytorch. - -# heating_profile: -# description: "Analyze heating profile characteristics" -# class: HeatingProfileAnalyzer -# file: heating_analyzer.py -# cli_flag: "--heating" -# default_config: {} -# outputs: -# - heating_rate_avg: "Average heating rate" -# - overshoot_celsius: "Temperature overshoot" -# - time_at_target: "Time at target temperature" -# requirements: -# - temperature_logs.parquet - diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/config_loader.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/config_loader.py deleted file mode 100644 index 5d6163017..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/config_loader.py +++ /dev/null @@ -1,327 +0,0 @@ -#!/usr/bin/env python3 -""" -Configuration Loader for A-Lab Pipeline - -Loads configuration with priority: -1. Environment variables (highest priority) -2. YAML config files -3. Code defaults (fallback) - -Environment Variables: - ALAB_MONGO_URI - MongoDB connection URI - ALAB_MONGO_DB - MongoDB database name - ALAB_MONGO_COLLECTION - MongoDB collection name - ALAB_S3_BUCKET - S3 bucket for uploads - ALAB_S3_PREFIX - S3 prefix path - ALAB_MP_API_KEY - Materials Project API key (for XRD analysis) - -Example: - export ALAB_MONGO_URI="mongodb://production:27017/" - export ALAB_S3_BUCKET="my-custom-bucket" -""" - -import os -import yaml -from pathlib import Path -from typing import Dict, Any, Optional -from functools import lru_cache - -# Config file paths -CONFIG_DIR = Path(__file__).parent -DEFAULTS_PATH = CONFIG_DIR / "defaults.yaml" - - -class ConfigLoader: - """Load configuration from env vars -> YAML -> defaults""" - - def __init__(self): - self._defaults = None - self._load_defaults() - - def _load_defaults(self): - """Load defaults.yaml file""" - if DEFAULTS_PATH.exists(): - with open(DEFAULTS_PATH, 'r') as f: - self._defaults = yaml.safe_load(f) or {} - else: - self._defaults = {} - - def _get_env_or_yaml(self, env_var: str, yaml_path: list, default: Any = None) -> Any: - """ - Get value from env var first, then YAML, then default. - - Args: - env_var: Environment variable name - yaml_path: Path through YAML dict (e.g., ['mongodb', 'uri']) - default: Fallback default value - - Returns: - Configuration value - """ - # 1. Try environment variable - env_value = os.getenv(env_var) - if env_value is not None: - return env_value - - # 2. Try YAML config - value = self._defaults - for key in yaml_path: - if isinstance(value, dict) and key in value: - value = value[key] - else: - value = None - break - - if value is not None: - return value - - # 3. Use default - return default - - # MongoDB configuration - @property - def mongo_uri(self) -> str: - """MongoDB connection URI""" - return self._get_env_or_yaml( - 'ALAB_MONGO_URI', - ['mongodb', 'uri'], - 'mongodb://localhost:27017/' - ) - - @property - def mongo_db(self) -> str: - """MongoDB database name""" - return self._get_env_or_yaml( - 'ALAB_MONGO_DB', - ['mongodb', 'database'], - 'temporary' - ) - - @property - def mongo_collection(self) -> str: - """MongoDB collection name""" - return self._get_env_or_yaml( - 'ALAB_MONGO_COLLECTION', - ['mongodb', 'collection'], - 'release' - ) - - # Parquet configuration - @property - def parquet_skip_temp_logs(self) -> bool: - """Skip temperature logs during transformation""" - value = self._get_env_or_yaml( - 'ALAB_SKIP_TEMP_LOGS', - ['parquet', 'skip_temperature_logs'], - False - ) - return self._to_bool(value) - - @property - def parquet_skip_xrd_points(self) -> bool: - """Skip XRD data points during transformation""" - value = self._get_env_or_yaml( - 'ALAB_SKIP_XRD_POINTS', - ['parquet', 'skip_xrd_points'], - False - ) - return self._to_bool(value) - - @property - def parquet_skip_workflow_tasks(self) -> bool: - """Skip workflow tasks during transformation""" - value = self._get_env_or_yaml( - 'ALAB_SKIP_WORKFLOW_TASKS', - ['parquet', 'skip_workflow_tasks'], - False - ) - return self._to_bool(value) - - @property - def parquet_compression(self) -> str: - """Parquet compression type""" - return self._get_env_or_yaml( - 'ALAB_PARQUET_COMPRESSION', - ['parquet', 'compression'], - 'snappy' - ) - - @property - def parquet_engine(self) -> str: - """Parquet engine""" - return self._get_env_or_yaml( - 'ALAB_PARQUET_ENGINE', - ['parquet', 'engine'], - 'pyarrow' - ) - - # S3 Upload configuration - @property - def s3_bucket(self) -> str: - """S3 bucket for uploads""" - return self._get_env_or_yaml( - 'ALAB_S3_BUCKET', - ['upload', 's3_bucket'], - 'materialsproject-contribs' - ) - - @property - def s3_prefix(self) -> str: - """S3 prefix path""" - return self._get_env_or_yaml( - 'ALAB_S3_PREFIX', - ['upload', 's3_prefix'], - 'alab_synthesis' - ) - - @property - def s3_exclude_large_files(self) -> bool: - """Exclude large files from S3 upload""" - value = self._get_env_or_yaml( - 'ALAB_S3_EXCLUDE_LARGE', - ['upload', 'exclude_large_files'], - True - ) - return self._to_bool(value) - - @property - def s3_large_file_threshold_mb(self) -> int: - """Large file threshold in MB""" - value = self._get_env_or_yaml( - 'ALAB_S3_LARGE_THRESHOLD_MB', - ['upload', 'large_file_threshold_mb'], - 50 - ) - return int(value) - - # Materials Project API - @property - def mp_api_key(self) -> Optional[str]: - """Materials Project API key (for XRD analysis)""" - return os.getenv('ALAB_MP_API_KEY') or os.getenv('MP_API_KEY') - - # Utility methods - @staticmethod - def _to_bool(value: Any) -> bool: - """Convert various types to boolean""" - if isinstance(value, bool): - return value - if isinstance(value, str): - return value.lower() in ('true', '1', 'yes', 'on') - return bool(value) - - def get_all(self) -> Dict[str, Any]: - """Get all configuration as dict""" - return { - 'mongodb': { - 'uri': self.mongo_uri, - 'database': self.mongo_db, - 'collection': self.mongo_collection, - }, - 'parquet': { - 'skip_temperature_logs': self.parquet_skip_temp_logs, - 'skip_xrd_points': self.parquet_skip_xrd_points, - 'skip_workflow_tasks': self.parquet_skip_workflow_tasks, - 'compression': self.parquet_compression, - 'engine': self.parquet_engine, - }, - 's3': { - 'bucket': self.s3_bucket, - 'prefix': self.s3_prefix, - 'exclude_large_files': self.s3_exclude_large_files, - 'large_file_threshold_mb': self.s3_large_file_threshold_mb, - }, - 'mp_api_key': self.mp_api_key, - } - - -# Singleton instance -@lru_cache(maxsize=1) -def get_config() -> ConfigLoader: - """Get cached configuration loader instance""" - return ConfigLoader() - - -# Convenience functions for direct access -def get_mongo_uri() -> str: - """Get MongoDB URI from env > yaml > default""" - return get_config().mongo_uri - - -def get_mongo_db() -> str: - """Get MongoDB database from env > yaml > default""" - return get_config().mongo_db - - -def get_mongo_collection() -> str: - """Get MongoDB collection from env > yaml > default""" - return get_config().mongo_collection - - -def get_s3_bucket() -> str: - """Get S3 bucket from env > yaml > default""" - return get_config().s3_bucket - - -def get_s3_prefix() -> str: - """Get S3 prefix from env > yaml > default""" - return get_config().s3_prefix - - -if __name__ == '__main__': - # Display current configuration - config = get_config() - - print("=" * 60) - print("A-Lab Pipeline Configuration") - print("=" * 60) - print("\nConfiguration Priority: ENV → YAML → Defaults\n") - - print("MongoDB:") - print(f" URI: {config.mongo_uri}") - print(f" Database: {config.mongo_db}") - print(f" Collection: {config.mongo_collection}") - - print("\nParquet:") - print(f" Skip temp logs: {config.parquet_skip_temp_logs}") - print(f" Skip XRD points: {config.parquet_skip_xrd_points}") - print(f" Skip workflow tasks: {config.parquet_skip_workflow_tasks}") - print(f" Compression: {config.parquet_compression}") - print(f" Engine: {config.parquet_engine}") - - print("\nS3 Upload:") - print(f" Bucket: {config.s3_bucket}") - print(f" Prefix: {config.s3_prefix}") - print(f" Exclude large files: {config.s3_exclude_large_files}") - print(f" Large file threshold: {config.s3_large_file_threshold_mb} MB") - - print("\nMaterials Project:") - mp_key = config.mp_api_key - if mp_key: - print(f" API Key: {mp_key[:10]}... (found)") - else: - print(f" API Key: Not set") - - print("\n" + "=" * 60) - - # Show which values are from env - print("\nEnvironment Variables Detected:") - env_vars = [ - 'ALAB_MONGO_URI', 'ALAB_MONGO_DB', 'ALAB_MONGO_COLLECTION', - 'ALAB_S3_BUCKET', 'ALAB_S3_PREFIX', - 'ALAB_SKIP_TEMP_LOGS', 'ALAB_SKIP_XRD_POINTS', - 'ALAB_MP_API_KEY', 'MP_API_KEY' - ] - - found_envs = [] - for var in env_vars: - if os.getenv(var): - found_envs.append(f" ✓ {var}") - - if found_envs: - print("\n".join(found_envs)) - else: - print(" (none) - using YAML/defaults") - - print("\n" + "=" * 60) - diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/defaults.yaml b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/defaults.yaml deleted file mode 100644 index da34782ad..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/defaults.yaml +++ /dev/null @@ -1,94 +0,0 @@ -# ============================================================================= -# A-Lab Pipeline Default Configuration -# ============================================================================= -# This file contains global defaults for the pipeline. -# Product-specific configs can override these values. -# -# Configuration Priority: -# 1. Environment variables (highest - for deployment/secrets) -# 2. This YAML file (middle - for general config) -# 3. Code defaults (lowest - fallback only) -# -# Environment Variables (see .env.example for full list): -# ALAB_MONGO_URI, ALAB_MONGO_DB, ALAB_MONGO_COLLECTION -# ALAB_S3_BUCKET, ALAB_S3_PREFIX -# ALAB_MP_API_KEY -# -# To customize: Either set env vars OR copy values to your product's config.yaml -# ============================================================================= - -# Pipeline version -version: '2.0' - -# MongoDB connection (used by mongodb_to_parquet.py) -# Override with: ALAB_MONGO_URI, ALAB_MONGO_DB, ALAB_MONGO_COLLECTION -mongodb: - uri: 'mongodb://localhost:27017/' - database: 'temporary' - collection: 'release' - -# Parquet transformation defaults -# Override with: ALAB_SKIP_TEMP_LOGS, ALAB_SKIP_XRD_POINTS, etc. -parquet: - # Skip large arrays by default for faster processing - skip_temperature_logs: false - skip_xrd_points: false - skip_workflow_tasks: false - # Compression settings - compression: 'snappy' - engine: 'pyarrow' - -# Default experiment filter (can be overridden per product) -experiment_filter: - # Status filter (completed experiments by default) - status: - - completed - # Require XRD data by default - has_xrd: true - # No date range by default (all dates) - date_range: null - -# Analysis defaults -analyses: - # XRD analysis settings - xrd_dara: - enabled: true - config: - wmin: 10 - wmax: 80 - save_viz: false - - # Powder statistics (always fast, recommend enabling) - powder_statistics: - enabled: true - config: {} - -# Upload settings -# Override with: ALAB_S3_BUCKET, ALAB_S3_PREFIX, ALAB_S3_EXCLUDE_LARGE -upload: - # S3 bucket for OpenData uploads - s3_bucket: 'materialsproject-contribs' - s3_prefix: 'alab_synthesis' - # Skip large files by default - exclude_large_files: true - large_file_threshold_mb: 50 - -# Diagram generation -diagram: - # Generate diagram after parquet creation - auto_generate: true - # Output formats - formats: - - summary # Markdown summary - - mermaid # ERD diagram - # Output filename (relative to product directory) - output_file: 'SCHEMA_DIAGRAM.md' - -# Validation settings -validation: - # Sample size for row-by-row validation (full validation can be slow) - sample_size: 100 - # Fail pipeline on validation errors - fail_on_error: false - # Show first N errors - max_errors_shown: 10 diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/filters.yaml b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/filters.yaml deleted file mode 100644 index 34de32236..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/config/filters.yaml +++ /dev/null @@ -1,97 +0,0 @@ -# ============================================================================= -# A-Lab Experiment Filter Presets -# ============================================================================= -# Predefined filters for common experiment selections. -# Use these in your product config or as templates. -# -# Usage in product config.yaml: -# experiment_filter: -# preset: "completed_with_xrd" # Use a preset -# OR -# experiment_filter: -# types: [NSC, Na] # Custom filter (types auto-discovered from MongoDB) -# ============================================================================= - -# NOTE: Experiment types are AUTO-DISCOVERED from MongoDB -# The pipeline discovers available types by querying MongoDB at runtime. -# Run this to see available types: -# python data/tools/analyze_mongodb.py -# -# Example types (discovered from A-Lab): -# NSC, Na, PG, MINES, TRI, etc. -# These change based on what's actually in your MongoDB database. - -# Valid experiment statuses (standardized) -statuses: - - completed # All tasks finished successfully - - error # At least one task failed - - active # Currently running - - unknown # Status could not be determined - -# ============================================================================= -# Filter Presets (Generic) -# ============================================================================= -# Generic filter presets that work regardless of experiment types in MongoDB. -# Type-specific filters are auto-generated during product creation. - -presets: - # All completed experiments with XRD data - completed_with_xrd: - description: 'Completed experiments that have XRD measurements' - status: - - completed - has_xrd: true - - # All completed experiments (with or without XRD) - all_completed: - description: 'All completed experiments' - status: - - completed - has_xrd: null # Don't filter by XRD - - # Failed experiments (for debugging) - failed: - description: 'Experiments with errors (for debugging)' - status: - - error - has_xrd: null - - # Everything (no filtering) - all: - description: 'All experiments (no filtering)' - # No filters applied -# ============================================================================= -# Type-Specific Filters (Auto-Generated) -# ============================================================================= -# Filters for specific experiment types (NSC, Na, PG, etc.) are auto-generated -# during product creation based on what's discovered in MongoDB. -# -# The pipeline will: -# 1. Query MongoDB for available experiment types -# 2. Show hierarchical selection (root types and subgroups) -# 3. Generate appropriate filters based on your selection -# -# You don't need to define type-specific presets here. -# ============================================================================= -# Custom Filters -# ============================================================================= -# Add your own generic filter presets here. -# For type-specific filters, use the interactive product creation which will -# auto-discover available types from MongoDB. -# -# Generic filter template: -# -# my_custom_filter: -# description: "What this filter selects" -# status: [completed] # Optional: status filter -# has_xrd: true # Optional: XRD requirement -# has_sem: true # Optional: SEM requirement (future) -# date_range: # Optional: date range -# start: "2024-01-01" -# end: "2024-12-31" -# -# NOTE: To filter by specific experiment types (e.g., NSC, Na): -# - Use interactive product creation (./run_product_pipeline.sh create) -# - Types are discovered from MongoDB at runtime -# - Or specify directly in product config.yaml: types: [NSC, Na] - diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/mongodb_to_parquet.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/mongodb_to_parquet.py deleted file mode 100644 index 92bdd7611..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/mongodb_to_parquet.py +++ /dev/null @@ -1,591 +0,0 @@ -#!/usr/bin/env python3 -""" -MongoDB to Parquet Data Pipeline -Transforms A-Lab experiment data from MongoDB to Parquet files for dashboard - -Output Structure (Consolidated): -- experiments.parquet: Main table with ALL 1:1 data merged (~45 columns) -- experiment_elements.parquet: Elements per experiment (1:N) -- powder_doses.parquet: Individual powder doses (1:N) -- temperature_logs.parquet: Temperature readings (1:N, optional) -- xrd_data_points.parquet: Raw XRD patterns (1:N, optional) -- workflow_tasks.parquet: Task execution history (1:N, optional) -- xrd_refinements.parquet: Analysis results (from DARA) -- xrd_phases.parquet: Identified phases (from DARA) -""" - -import sys -import os -import logging -from pathlib import Path -from datetime import datetime -from typing import Dict, List, Optional, Any -import argparse - -import pandas as pd -import numpy as np -from pymongo import MongoClient -from tqdm import tqdm - -# Import config loader -sys.path.insert(0, str(Path(__file__).parent / "config")) -from config_loader import get_config - -# Configure logging -logging.basicConfig( - level=logging.INFO, - format='%(asctime)s - %(levelname)s - %(message)s', - handlers=[ - logging.FileHandler('parquet_migration.log'), - logging.StreamHandler() - ] -) -logger = logging.getLogger(__name__) - -# Load configuration (env vars > yaml > defaults) -config = get_config() - -# Output directory for parquet files -OUTPUT_DIR = Path(__file__).parent / "parquet" - - -class MongoToParquetTransformer: - """Transform MongoDB experiments to consolidated Parquet files""" - - def __init__(self, mongo_uri: str = None, mongo_db: str = None, - mongo_collection: str = None, output_dir: Path = OUTPUT_DIR): - # Use config if not explicitly provided - self.mongo_uri = mongo_uri or config.mongo_uri - self.mongo_db = mongo_db or config.mongo_db - self.mongo_collection = mongo_collection or config.mongo_collection - - self.output_dir = Path(output_dir) - self.output_dir.mkdir(parents=True, exist_ok=True) - - # Connect to MongoDB - self.client = MongoClient(self.mongo_uri) - self.db = self.client[self.mongo_db] - self.collection = self.db[self.mongo_collection] - - logger.info(f"Connected to MongoDB: {mongo_uri}") - logger.info(f"Output directory: {self.output_dir}") - - # Counter for skipped experiments - self.skipped_no_sampleid = 0 - - # Data containers - consolidated structure - # Main experiments table (merged 1:1 data) - self.experiments = [] - - # 1:N tables (must stay separate) - self.experiment_elements = [] - self.powder_doses = [] - self.temperature_logs = [] - self.xrd_data_points = [] - self.workflow_tasks = [] - - def _build_filter_query(self, experiment_filter: Dict) -> Dict: - """Build MongoDB query from filter criteria""" - query = {} - - if 'types' in experiment_filter and experiment_filter['types']: - # Match experiments starting with any of the prefixes - patterns = [f"^{t}_" for t in experiment_filter['types']] - query["name"] = {"$regex": "|".join(patterns)} - - if 'status' in experiment_filter and experiment_filter['status']: - query["status"] = {"$in": experiment_filter['status']} - - if 'has_xrd' in experiment_filter: - if experiment_filter['has_xrd']: - query["metadata.diffraction_results.sampleid_in_aeris"] = {"$exists": True, "$ne": None} - else: - query["$or"] = [ - {"metadata.diffraction_results.sampleid_in_aeris": {"$exists": False}}, - {"metadata.diffraction_results.sampleid_in_aeris": None} - ] - - if 'date_range' in experiment_filter: - date_query = {} - if 'start' in experiment_filter['date_range']: - date_query["$gte"] = experiment_filter['date_range']['start'] - if 'end' in experiment_filter['date_range']: - date_query["$lte"] = experiment_filter['date_range']['end'] - if date_query: - query["last_updated"] = date_query - - if 'experiment_names' in experiment_filter and experiment_filter['experiment_names']: - query["name"] = {"$in": experiment_filter['experiment_names']} - - return query - - def transform_all(self, limit: Optional[int] = None, - skip_temp_logs: bool = False, - skip_xrd_points: bool = False, - skip_workflow_tasks: bool = False, - experiment_filter: Optional[Dict] = None): - """Transform all experiments to dataframes - - Args: - limit: Maximum number of experiments to process - skip_temp_logs: Skip temperature log data (large arrays) - skip_xrd_points: Skip XRD data points (very large arrays) - skip_workflow_tasks: Skip workflow task history - experiment_filter: Filter criteria for selecting experiments - - types: List of experiment type prefixes (e.g., ['NSC', 'Na']) - - status: List of statuses (e.g., ['completed']) - - has_xrd: Boolean, whether XRD data is required - - date_range: Dict with 'start' and/or 'end' dates - """ - - query = self._build_filter_query(experiment_filter) if experiment_filter else {} - total = self.collection.count_documents(query) - - if limit: - total = min(limit, total) - - logger.info(f"Processing {total} experiments...") - if experiment_filter: - logger.info(f"Applied filter: {experiment_filter}") - - cursor = self.collection.find(query).limit(limit) if limit else self.collection.find(query) - - processed = 0 - for doc in tqdm(cursor, total=total, desc="Transforming experiments"): - try: - if self._transform_experiment(doc, skip_temp_logs, skip_xrd_points, skip_workflow_tasks): - processed += 1 - except Exception as e: - logger.error(f"Error transforming {doc.get('name', 'unknown')}: {e}") - continue - - # Log filtering statistics - logger.info("-" * 40) - logger.info("Filtering Summary:") - logger.info(f" Total experiments in MongoDB: {total}") - logger.info(f" Skipped (no sampleid_in_aeris): {self.skipped_no_sampleid}") - logger.info(f" Processed: {processed}") - logger.info("-" * 40) - logger.info("Transformation complete. Creating dataframes...") - self._save_all_parquet_files() - - def _transform_experiment(self, doc: Dict, skip_temp_logs: bool, - skip_xrd_points: bool, skip_workflow_tasks: bool) -> bool: - """Transform a single experiment document into consolidated format - - Returns: - bool: True if experiment was processed, False if skipped - """ - - exp_name = doc.get('name', '') - metadata = doc.get('metadata', {}) - - # Filter: Skip experiments without sampleid_in_aeris (indicates no XRD performed) - xrd = metadata.get('diffraction_results', {}) - sampleid = xrd.get('sampleid_in_aeris') if xrd else None - if not sampleid: - logger.debug(f"Skipping {exp_name}: No sampleid_in_aeris (no XRD)") - self.skipped_no_sampleid += 1 - return False - - exp_id = str(doc['_id']) - tasks = doc.get('tasks', []) or [] - - # Extract hierarchical type from name: NSC_249_001 -> root=NSC, subgroup=NSC_249 - parts = exp_name.split('_') if '_' in exp_name else [exp_name] - exp_type = parts[0] if parts else 'Unknown' - exp_subgroup = '_'.join(parts[:2]) if len(parts) >= 2 else None - - # ========================================== - # Build consolidated experiment record - # Merging: experiments + heating + recovery + xrd + finalization + dosing - # ========================================== - - experiment_record = { - # Core experiment fields - 'experiment_id': exp_id, - 'name': exp_name, - 'experiment_type': exp_type, - 'experiment_subgroup': exp_subgroup, - 'target_formula': metadata.get('target', ''), - 'last_updated': doc.get('last_updated'), - 'status': self._get_experiment_status(tasks), - 'notes': None, - } - - # --- Heating session fields (prefix: heating_) --- - heating = metadata.get('heating_results', {}) - heating_method = self._determine_heating_method(tasks, heating) - - experiment_record.update({ - 'heating_method': heating_method, - 'heating_temperature': heating.get('heating_temperature'), - 'heating_time': heating.get('heating_time'), - 'heating_cooling_rate': heating.get('cooling_rate'), - 'heating_atmosphere': heating.get('atmosphere'), - 'heating_flow_rate_ml_min': heating.get('flow_rate'), - 'heating_low_temp_calcination': heating.get('low_temperature_calcination', False), - }) - - # --- Powder recovery fields (prefix: recovery_) --- - recovery = metadata.get('recoverpowder_results', {}) - dosing = metadata.get('powderdosing_results', {}) - - # Calculate total dosed mass - total_dosed_mg = self._calculate_total_dosed_mass(dosing) - collected_weight = recovery.get('weight_collected') - yield_pct = None - if total_dosed_mg and total_dosed_mg > 0 and collected_weight: - yield_pct = (collected_weight / total_dosed_mg) * 100 - - experiment_record.update({ - 'recovery_total_dosed_mass_mg': total_dosed_mg if total_dosed_mg > 0 else None, - 'recovery_weight_collected_mg': collected_weight, - 'recovery_yield_percent': yield_pct, - 'recovery_initial_crucible_weight_mg': recovery.get('initial_crucible_weight'), - 'recovery_failure_classification': recovery.get('failure_classification'), - }) - - # --- XRD measurement fields (prefix: xrd_) --- - experiment_record.update({ - 'xrd_sampleid_in_aeris': xrd.get('sampleid_in_aeris'), - 'xrd_holder_index': xrd.get('xrd_holder_index'), - 'xrd_total_mass_dispensed_mg': xrd.get('total_mass_dispensed_mg'), - 'xrd_met_target_mass': xrd.get('met_target_mass'), - }) - - # --- Sample finalization fields (prefix: finalization_) --- - ending = metadata.get('ending_results', {}) - experiment_record.update({ - 'finalization_decoded_sample_id': ending.get('decoded_sample_id'), - 'finalization_successful_labeling': ending.get('successful_labeling'), - 'finalization_storage_location': self._get_storage_location(tasks), - }) - - # --- Dosing session fields (prefix: dosing_) --- - if dosing: - experiment_record.update({ - 'dosing_crucible_position': dosing.get('CruciblePosition'), - 'dosing_crucible_sub_rack': dosing.get('CrucibleSubRack'), - 'dosing_mixing_pot_position': dosing.get('MixingPotPosition'), - 'dosing_ethanol_dispense_volume': dosing.get('EthanolDispenseVolume'), - 'dosing_target_transfer_volume': dosing.get('TargetTransferVolume'), - 'dosing_actual_transfer_mass': dosing.get('ActualTransferMass'), - 'dosing_dac_duration': dosing.get('DACDuration'), - 'dosing_dac_speed': dosing.get('DACSpeed'), - 'dosing_actual_heat_duration': dosing.get('ActualHeatDuration'), - 'dosing_end_reason': dosing.get('EndReason'), - }) - - self.experiments.append(experiment_record) - - # ========================================== - # 1:N Tables (must stay separate) - # ========================================== - - # Experiment elements - elements = metadata.get('elements_present', []) or [] - for element in elements: - self.experiment_elements.append({ - 'experiment_id': exp_id, - 'element_symbol': element, - 'target_atomic_percent': None - }) - - # Powder doses (nested structure) - powders = dosing.get('Powders', []) or [] - for powder in powders: - powder_name = powder.get('PowderName', '') - target_mass = powder.get('TargetMass', 0) - - doses = powder.get('Doses', []) or [] - for idx, dose in enumerate(doses): - actual_mass = dose.get('Mass', 0) - accuracy = ((actual_mass / target_mass * 100) - if target_mass > 0 else None) - - self.powder_doses.append({ - 'experiment_id': exp_id, - 'powder_name': powder_name, - 'target_mass': target_mass, - 'actual_mass': actual_mass, - 'accuracy_percent': accuracy, - 'dose_sequence': idx, - 'head_position': dose.get('HeadPosition'), - 'dose_timestamp': dose.get('TimeStamp') - }) - - # Temperature logs (large arrays, optional) - if not skip_temp_logs and heating: - temp_log = heating.get('temperature_log', {}) or {} - times = temp_log.get('time_minutes', []) or [] - temps = temp_log.get('temperature_celsius', []) or [] - - for idx, (time_min, temp_c) in enumerate(zip(times, temps)): - self.temperature_logs.append({ - 'experiment_id': exp_id, - 'sequence_number': idx, - 'time_minutes': time_min, - 'temperature_celsius': temp_c - }) - - # XRD data points (HUGE arrays, optional) - if not skip_xrd_points and xrd: - twotheta = xrd.get('twotheta', []) or [] - counts = xrd.get('counts', []) or [] - - for idx, (theta, count) in enumerate(zip(twotheta, counts)): - self.xrd_data_points.append({ - 'experiment_id': exp_id, - 'point_index': idx, - 'twotheta': theta, - 'counts': count - }) - - # Workflow tasks (optional) - if not skip_workflow_tasks: - for task in tasks: - self.workflow_tasks.append({ - 'experiment_id': exp_id, - 'task_id': str(task.get('task_id', '')), - 'task_type': task.get('type', ''), - 'status': task.get('status', ''), - 'created_at': task.get('created_at'), - 'started_at': task.get('started_at'), - 'completed_at': task.get('completed_at') - }) - - return True - - def _determine_heating_method(self, tasks: List[Dict], heating: Dict) -> str: - """Determine heating method from task types""" - heating_method = 'none' - for task in tasks: - task_type = task.get('type', '') - if task_type == 'HeatingWithAtmosphere': - heating_method = 'atmosphere' - params = task.get('parameters', {}) - if params.get('atmosphere'): - heating['atmosphere'] = params.get('atmosphere') - heating['flow_rate'] = params.get('flow_rate') - break - elif task_type == 'ManualHeating': - heating_method = 'manual' - break - elif task_type == 'Heating': - heating_method = 'standard' - break - return heating_method - - def _calculate_total_dosed_mass(self, dosing: Dict) -> float: - """Calculate total dosed mass from powder doses (grams -> mg)""" - total_dosed_mg = 0 - powders = dosing.get('Powders', []) or [] - for powder in powders: - doses = powder.get('Doses', []) or [] - for dose in doses: - total_dosed_mg += (dose.get('Mass', 0) or 0) * 1000 # g to mg - return total_dosed_mg - - def _get_experiment_status(self, tasks: List[Dict]) -> str: - """Determine experiment status from tasks""" - if not tasks or tasks is None: - return 'unknown' - - statuses = [t.get('status', '') for t in tasks] - - if all(s == 'COMPLETED' for s in statuses): - return 'completed' - elif any(s == 'ERROR' for s in statuses): - return 'error' - elif any(s in ['RUNNING', 'WAITING'] for s in statuses): - return 'active' - else: - return 'unknown' - - def _get_storage_location(self, tasks: List[Dict]) -> Optional[str]: - """Extract final storage location from Ending task""" - if not tasks: - return None - - for task in tasks: - if task.get('type') == 'Ending': - subtasks = task.get('subtasks', []) - if subtasks: - last_subtask = subtasks[-1] - params = last_subtask.get('parameters', {}) - return params.get('destination') - return None - - def _save_all_parquet_files(self): - """Convert all lists to dataframes and save as parquet""" - - datasets = { - 'experiments': self.experiments, - 'experiment_elements': self.experiment_elements, - 'powder_doses': self.powder_doses, - 'temperature_logs': self.temperature_logs, - 'xrd_data_points': self.xrd_data_points, - 'workflow_tasks': self.workflow_tasks, - } - - for name, data in datasets.items(): - if not data: - logger.warning(f"No data for {name}, skipping...") - continue - - df = pd.DataFrame(data) - output_path = self.output_dir / f"{name}.parquet" - - # Optimize data types for better compression - df = self._optimize_dtypes(df) - - # Save with compression (from config) - df.to_parquet( - output_path, - engine=config.parquet_engine, - compression=config.parquet_compression, - index=False - ) - - file_size_mb = output_path.stat().st_size / (1024 * 1024) - logger.info(f"✓ Saved {name}: {len(df):,} rows, {len(df.columns)} cols, {file_size_mb:.2f} MB") - - # Save metadata - self._save_metadata(datasets) - - def _optimize_dtypes(self, df: pd.DataFrame) -> pd.DataFrame: - """Optimize dataframe dtypes for smaller file size""" - - for col in df.columns: - col_type = df[col].dtype - - # Downcast integers - if col_type == 'int64': - df[col] = pd.to_numeric(df[col], downcast='integer') - - # Downcast floats - elif col_type == 'float64': - df[col] = pd.to_numeric(df[col], downcast='float') - - # Convert to category if low cardinality - elif col_type == 'object': - num_unique = df[col].nunique() - num_total = len(df) - if num_unique / num_total < 0.5: # Less than 50% unique - df[col] = df[col].astype('category') - - return df - - def _save_metadata(self, datasets: Dict[str, List]): - """Save migration metadata""" - - # Get column info for experiments table - exp_columns = list(pd.DataFrame(datasets['experiments']).columns) if datasets['experiments'] else [] - - metadata = { - 'migration_date': datetime.now().isoformat(), - 'version': '2.0', # Consolidated schema version - 'source': { - 'type': 'mongodb', - 'uri': self.mongo_uri, - 'database': self.mongo_db, - 'collection': self.mongo_collection - }, - 'output_directory': str(self.output_dir), - 'schema': { - 'description': 'Consolidated schema - all 1:1 tables merged into experiments.parquet', - 'experiments_columns': len(exp_columns), - 'merged_tables': ['experiments', 'heating_sessions', 'powder_recovery', - 'xrd_measurements', 'sample_finalization', 'dosing_sessions'] - }, - 'datasets': { - name: { - 'rows': len(data), - 'file': f"{name}.parquet" - } - for name, data in datasets.items() if data - } - } - - import json - metadata_path = self.output_dir / 'metadata.json' - with open(metadata_path, 'w') as f: - json.dump(metadata, f, indent=2) - - logger.info(f"✓ Saved metadata: {metadata_path}") - - def close(self): - """Close MongoDB connection""" - self.client.close() - - -def main(): - parser = argparse.ArgumentParser( - description='Transform MongoDB experiments to consolidated Parquet files' - ) - parser.add_argument( - '--output-dir', '-o', - type=Path, - default=OUTPUT_DIR, - help='Output directory for parquet files' - ) - parser.add_argument( - '--mongo-uri', '-m', - default=None, - help=f'MongoDB connection URI (default: from config, currently {config.mongo_uri})' - ) - parser.add_argument( - '--limit', '-l', - type=int, - help='Limit number of experiments to process (for testing)' - ) - parser.add_argument( - '--skip-temp-logs', - action='store_true', - help='Skip temperature log data (saves ~500k rows)' - ) - parser.add_argument( - '--skip-xrd-points', - action='store_true', - help='Skip XRD data points (saves ~4.7M rows)' - ) - parser.add_argument( - '--skip-workflow-tasks', - action='store_true', - help='Skip workflow task history (saves ~3k rows)' - ) - - args = parser.parse_args() - - logger.info("="*60) - logger.info("MongoDB to Parquet Transformation Pipeline (v2 - Consolidated)") - logger.info("="*60) - - try: - transformer = MongoToParquetTransformer( - mongo_uri=args.mongo_uri, - output_dir=args.output_dir - ) - - transformer.transform_all( - limit=args.limit, - skip_temp_logs=args.skip_temp_logs, - skip_xrd_points=args.skip_xrd_points, - skip_workflow_tasks=args.skip_workflow_tasks - ) - - logger.info("="*60) - logger.info("✓ Transformation complete!") - logger.info(f"✓ Parquet files saved to: {args.output_dir}") - logger.info("="*60) - - except Exception as e: - logger.error(f"Pipeline failed: {e}", exc_info=True) - sys.exit(1) - finally: - transformer.close() - - -if __name__ == '__main__': - main() diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/PIPELINE_ARCHITECTURE.md b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/PIPELINE_ARCHITECTURE.md deleted file mode 100644 index 024b58a13..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/PIPELINE_ARCHITECTURE.md +++ /dev/null @@ -1,472 +0,0 @@ -# A-Lab Pipeline Architecture v2 - -## Overview - -The A-Lab pipeline is a product-based data processing system that: - -1. Filters experiments by configurable criteria -2. Transforms MongoDB data to Parquet (uses `mongodb_to_parquet.py` primitive) -3. Runs **auto-discovered** analyses (XRD, powder statistics, custom) -4. Validates data with **auto-discovered** Pydantic schemas -5. Generates schema diagrams per product -6. Uploads parquet files directly to S3 OpenData - -### Key Features - -- **Auto-Discovery**: Schemas and analyses are discovered automatically from their directories -- **Extensible**: Add new data types or analyses by creating files, no code changes needed -- **Configuration-Driven**: Defaults and filters defined in YAML files - -## Quick Start - -### Create a New Data Product - -```bash -# Interactive product creation (uses default schema automatically) -python data/pipeline/run_pipeline.py create - -# Follow prompts to: -# - Name your product (e.g., "reaction_genome") -# - Select experiment groups hierarchically: -# • Root types (NSC, Na, PG, MINES, TRI) - selects ALL experiments -# • Subgroups (NSC_249, Na_123_A) - selects ONLY those experiments -# - Configure filters (has_xrd, status, date_range) -# - Choose analyses to run -# - Schema is loaded from data/products/default_schema.yaml (source of truth) -# - Optionally customize by removing unwanted fields -``` - -### Run the Pipeline - -```bash -# Dry run (default) -python data/pipeline/run_pipeline.py run --product reaction_genome - -# Live upload to S3 OpenData -python data/pipeline/run_pipeline.py run --product reaction_genome --upload - -# Run specific stages only -python data/pipeline/run_pipeline.py run --product reaction_genome --stages filter transform -``` - -### List Products - -```bash -python data/pipeline/run_pipeline.py list -``` - -### Regenerate Pydantic Schema - -```bash -# Regenerate schema.py from config.yaml (auto-syncs on every run) -python data/pipeline/run_pipeline.py regenerate --product reaction_genome -``` - -## Architecture - -``` -┌─────────────────────────────────────────────────────────────────┐ -│ Product Configuration │ -│ data/products/reaction_genome/config.yaml │ -│ │ -│ • Experiment filters (types, status, has_xrd) │ -│ • Product metadata (title, authors, description) │ -│ • Analysis pipeline (xrd_dara, powder_stats, etc.) │ -│ • Schema with units (heatingTemperature: degC) │ -└─────────────────────────────────────────────────────────────────┘ - ↓ -┌─────────────────────────────────────────────────────────────────┐ -│ Filter Experiments │ -│ Hierarchical Selection from MongoDB │ -│ │ -│ Discover: NSC (218) → NSC_249 (12), NSC_250 (8), ... │ -│ Na (215) → Na_123_A (5), Na_123_B (7), ... │ -│ Select: Root types OR specific subgroups │ -│ Apply filters: status=completed, has_xrd=true │ -│ Result: Filtered experiments → experiments.txt │ -└─────────────────────────────────────────────────────────────────┘ - ↓ -┌─────────────────────────────────────────────────────────────────┐ -│ Transform to Parquet │ -│ Create Product-Specific Filtered Data │ -│ │ -│ Each product maintains its own filtered parquet dataset: │ -│ • Data consistency with submitted version │ -│ • No cross-product contamination │ -│ • Schema version tracking per product │ -│ │ -│ Input: Filtered experiment list (from experiments.txt) │ -│ Output: products/{name}/parquet_data/*.parquet (filtered) │ -│ Schema: default_schema.yaml (new) OR config.yaml (submitted) │ -│ Tables: experiments, heating_sessions, xrd_measurements, etc. │ -└─────────────────────────────────────────────────────────────────┘ - ↓ -┌─────────────────────────────────────────────────────────────────┐ -│ Run Analysis Pipeline │ -│ Pluggable Analyzers │ -│ │ -│ • XRDAnalyzer → xrd_success, rwp, num_phases │ -│ • PowderStatisticsAnalyzer → avg_accuracy, total_doses │ -│ • Custom analyzers can be added │ -│ Output: products/reaction_genome/analysis_results/*.parquet │ -└─────────────────────────────────────────────────────────────────┘ - ↓ -┌─────────────────────────────────────────────────────────────────┐ -│ Validate with Pydantic │ -│ Auto-generated Pydantic + Schema Validation │ -│ │ -│ Pydantic-Based Validation Strategy: │ -│ • SchemaManager discovers all YAML schemas in schema/ │ -│ • Generates Pydantic model for each table (*_schema.py) │ -│ • SchemaValidator validates each parquet table row-by-row │ -│ • Leverages Pydantic's type checking and error messages │ -│ │ -│ 1. Auto-discover schemas: experiments, powder_doses, etc. │ -│ 2. Generate Pydantic models from YAML schemas │ -│ 3. Import Pydantic models dynamically at validation time │ -│ 4. Validate every row in every parquet table │ -│ 5. Report validation statistics and detailed errors │ -│ 6. Schema-agnostic: add new table → auto validates │ -└─────────────────────────────────────────────────────────────────┘ - ↓ -┌─────────────────────────────────────────────────────────────────┐ -│ Upload to S3 OpenData │ -│ │ -│ 1. Upload parquet files directly to S3 bucket │ -│ 2. Embed metadata in parquet schema │ -│ 3. Exclude very large files (temperature_logs, xrd_data_points) │ -│ 4. Available via S3 API and HTTP │ -└─────────────────────────────────────────────────────────────────┘ -``` - -## File Structure - -``` -A-Lab_Samples/ -├── data/ -│ ├── config/ # ⭐ Global configuration -│ │ ├── defaults.yaml # Pipeline defaults -│ │ ├── filters.yaml # Filter presets & experiment types -│ │ └── analyses.yaml # Analysis documentation -│ │ -│ ├── products/ # Data product definitions -│ │ ├── base_product.py # Product configuration system -│ │ ├── schema_manager.py # Auto-discovers Pydantic schemas -│ │ ├── schema_validator.py # Validates parquet vs schema -│ │ ├── schema/ # ⭐ AUTO-DISCOVERED schemas -│ │ │ ├── base.py # Shared utilities (ExcludeFromUpload) -│ │ │ ├── experiments.py # Main experiment schema -│ │ │ ├── powder_doses.py # Powder dosing schema -│ │ │ ├── xrd_*.py # XRD-related schemas -│ │ │ └── [your_schema.py] # ⭐ Add your own - auto-discovered! -│ │ │ -│ │ └── {product_name}/ # Product-specific directory -│ │ ├── config.yaml # Product configuration -│ │ ├── experiments.txt # Filtered experiment list -│ │ ├── parquet_data/ # Transformed data -│ │ ├── analysis_results/ # Analysis outputs -│ │ └── SCHEMA_DIAGRAM.md # Auto-generated diagram -│ │ -│ ├── pipeline/ -│ │ ├── run_pipeline.py # Main CLI (create, run, list) -│ │ ├── product_pipeline.py # Pipeline orchestration -│ │ ├── s3_uploader.py # S3 OpenData uploader -│ │ ├── archive/ # Deprecated API-based uploaders -│ │ └── pipeline_state.py # State tracking -│ │ -│ ├── analyses/ # ⭐ AUTO-DISCOVERED analysis plugins -│ │ ├── base_analyzer.py # Base class + built-in analyzers -│ │ └── [your_analyzer.py] # ⭐ Add your own - auto-discovered! -│ │ -│ ├── mongodb_to_parquet.py # Shared primitive: MongoDB → Parquet -│ │ -│ ├── xrd_creation/ # XRD analysis (DARA) -│ │ └── analyze_batch.py # Entry point: run_analysis.sh -│ │ -│ └── tools/ -│ ├── generate_diagram.py # Schema diagram generator -│ └── analyze_mongodb.py # Discover experiment types -│ -├── update_data.sh # MongoDB → Parquet (all data) -├── run_analysis.sh # XRD analysis runner -├── run_dashboard.sh # Dashboard (operates on parquet) -└── run_product_pipeline.sh # Product pipeline CLI wrapper -``` - -### Auto-Discovery Directories - -| Directory | What's Discovered | How to Extend | -| -------------------------- | ----------------- | ------------------------------------- | -| `data/products/schema/` | Pydantic schemas | Add `*.py` with BaseModel class | -| `data/analyses/` | Analysis plugins | Add `*.py` with BaseAnalyzer subclass | -| `data/config/filters.yaml` | Filter presets | Add YAML entry | - -## Schema Management (Auto-Discovered) - -Schemas are **auto-discovered** from `data/products/schema/`. Python-first approach! - -### Schema Directory Structure - -``` -data/products/schema/ -├── base.py # Shared utilities (ExcludeFromUpload) -├── experiments.py # Main experiment schema (~40 fields) -├── experiment_elements.py # Elements per experiment -├── powder_doses.py # Powder dosing data -├── temperature_logs.py # Temperature time series -├── xrd_data_points.py # XRD patterns -├── xrd_refinements.py # DARA analysis results -├── xrd_phases.py # Identified phases -└── [your_schema.py] # ⭐ Auto-discovered! -``` - -### Adding a New Schema - -```python -# data/products/schema/sem_data.py -"""SEM measurement schema""" - -from pydantic import BaseModel, Field - -class SEMData(BaseModel, extra="forbid"): - """SEM measurement data""" - - # Optional: specify table name (defaults to filename) - __schema_table__ = "sem_data" - - experiment_id: str = Field(description="Experiment ID") - image_count: int = Field(description="Number of SEM images") - morphology: str | None = Field(default=None, description="Classified morphology") -``` - -### List Available Schemas - -```bash -# Show all auto-discovered schemas -python data/products/schema_manager.py -``` - -### Marking Fields as Embargoed - -Use `ExcludeFromUpload` to prevent sensitive fields from being uploaded: - -```python -from .base import ExcludeFromUpload - -class MySchema(BaseModel): - public_field: str = Field(description="This is uploaded") - - # This field is NOT uploaded to S3 - sensitive_field: float | None = ExcludeFromUpload( - description="Embargoed until publication" - ) -``` - -### Schema Workflow - -``` -data/products/schema/*.py (Python-first, auto-discovered) - ↓ -SchemaManager loads all schemas automatically - ↓ -SchemaValidator validates parquet data - ↓ -Embargoed fields excluded from upload -``` - -⚠️ **Schemas are Python code** - edit directly, no YAML regeneration needed! - -```yaml -name: reaction_genome -version: 1.0 - -experiment_filter: - types: [NSC, Na] # Experiment prefixes - status: [completed] # Workflow status - has_xrd: true # Must have XRD data - -metadata: - title: 'A-Lab NASICON Synthesis Dataset' - authors: 'A-Lab Team, LBNL' - description: 'High-throughput synthesis data' - references: - - label: github - url: https://github.com/... - -analyses: - - name: xrd_dara - enabled: true - config: - wmin: 10 - wmax: 80 - - - name: powder_statistics - enabled: true - -schema: - heatingTemperature: - unit: degC - description: 'Synthesis temperature' - type: float - required: true - min: 0 - max: 2000 - - recoveryYield: - unit: '%' - description: 'Powder recovery percentage' - type: float - required: false - min: 0 - max: 100 -``` - -## Analysis Plugins (Auto-Discovered) - -Analyses are **auto-discovered** from `data/analyses/`. No registration needed! - -### Creating a Custom Analyzer - -```python -# data/analyses/my_analyzer.py - -from pathlib import Path -import pandas as pd -from base_analyzer import BaseAnalyzer - -class MyCustomAnalyzer(BaseAnalyzer): - """My custom analysis for A-Lab experiments""" - - # Class attributes for discovery (all optional except 'name') - name = "my_analysis" # Required: unique identifier - description = "Calculate my metrics" # Optional: shown in list - cli_flag = "--my-analysis" # Optional: for CLI usage - - def analyze(self, experiments_df: pd.DataFrame, parquet_dir: Path) -> pd.DataFrame: - """Run analysis on experiments""" - results = [] - - for _, exp in experiments_df.iterrows(): - results.append({ - 'experiment_name': exp['name'], - 'my_metric': self._calculate_metric(exp) - }) - - return pd.DataFrame(results) - - def get_output_schema(self): - return { - 'my_metric': { - 'type': 'float', - 'required': True, - 'description': 'My custom metric' - } - } - - def _calculate_metric(self, exp): - # Your analysis logic here - return 0.0 -``` - -### List Available Analyzers - -```bash -# Show all auto-discovered analyzers -python data/analyses/base_analyzer.py -``` - -### Enable in Product Config - -```yaml -# data/products/my_product/config.yaml -analyses: - - name: my_analysis - enabled: true - config: - param1: value1 -``` - -## S3 OpenData Integration - -### AWS Credentials Setup - -Create `.env` file at project root: - -``` -aws_access_key_id=your_key_here -aws_secret_access_key=your_secret_here -``` - -### Data Format - -Parquet files with embedded metadata: - -1. **Column metadata**: Units and descriptions in arrow schema -2. **Project metadata**: Embedded as schema metadata -3. **Efficient storage**: Columnar format for fast partial reads - -### Upload Process - -1. **Prepare parquet files**: Product-specific filtered data -2. **Embed metadata**: Add project info to schema -3. **Upload to S3**: Direct upload to `s3://materialsproject-contribs/alab_synthesis/{product}/` -4. **Exclude large files**: Temperature logs and XRD data points skipped by default - -## State Management - -Pipeline state tracked in `pipeline_runs.parquet`: - -- Run ID and timestamp -- Experiments processed -- Upload status -- Dry run flag - -View state: - -```python -from pipeline_state import PipelineStateManager -mgr = PipelineStateManager() -mgr.get_status_summary() -``` - -## Troubleshooting - -### "No experiments match filter" - -- Check filter criteria in config.yaml -- Use `python data/tools/analyze_mongodb.py` to explore available types -- Check `data/config/filters.yaml` for preset options - -### "Validation errors" - -- Check Pydantic schema matches your data -- Run `python data/products/schema_manager.py` to see loaded schemas -- Verify min/max ranges are appropriate - -### "S3 upload failed" - -- Verify AWS credentials in .env -- Check AWS access permissions -- Ensure boto3 is installed: `pip install boto3` -- Verify network connectivity to S3 - -### "Analysis not found" - -- Run `python data/analyses/base_analyzer.py` to see available analyzers -- Ensure analyzer file is in `data/analyses/` directory -- Check that analyzer class inherits from `BaseAnalyzer` -- Verify `name` class attribute is set - -### "Schema not found" - -- Run `python data/products/schema_manager.py` to see discovered schemas -- Ensure schema file is in `data/products/schema/` directory -- Check that schema class inherits from `pydantic.BaseModel` -- Verify file is not in SKIP_FILES list (base.py, **init**.py) - -## Best Practices - -1. **Start with dry runs**: Test pipeline before uploading -2. **Validate early**: Use Pydantic to catch issues -3. **Modular analyses**: Keep analyzers independent -4. **Version configs**: Track changes to product definitions diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/README.md b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/README.md deleted file mode 100644 index e696dd524..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/README.md +++ /dev/null @@ -1,267 +0,0 @@ -# A-Lab Data Pipeline - -Orchestrates data flow from MongoDB → Parquet → XRD Analysis → S3 OpenData. - -## 📚 Documentation - -- **[PARQUET_STRUCTURE.md](./PARQUET_STRUCTURE.md)** - Complete guide to parquet files (14 tables, relationships, usage) -- **[PIPELINE_ARCHITECTURE.md](./PIPELINE_ARCHITECTURE.md)** - Product-based pipeline design -- **[requirements.txt](./requirements.txt)** - Python dependencies - -## Quick Start - -```bash -# Create a new data product -python run_pipeline.py create - -# List available products -python run_pipeline.py list - -# Run pipeline (dry run) -python run_pipeline.py run --product reaction_genome - -# Run pipeline with live upload to S3 -python run_pipeline.py run --product reaction_genome --upload - -# Check status -python run_pipeline.py status --product reaction_genome -``` - -## Pipeline Phases - -| Phase | Script | Output | -| --------------- | ------------------ | --------------------------------------------------------------------- | -| 1. MongoDB Sync | `update_data.sh` | 14 parquet files (see [PARQUET_STRUCTURE.md](./PARQUET_STRUCTURE.md)) | -| 2. XRD Analysis | `analyze_batch.py` | `xrd_refinements.parquet`, `xrd_phases.parquet` | -| 3. S3 Upload | `s3_uploader.py` | Parquet files on s3://materialsproject-contribs | - -### Parquet Files Generated - -The pipeline creates **14 normalized parquet tables** (not one big file): - -| Category | File | Rows | Purpose | -| ----------------- | --------------------------- | ----- | ------------------------------------ | -| **Core** | experiments.parquet | 592 | Main experiment metadata | -| **Related (1:1)** | heating_sessions.parquet | 592 | Heating parameters | -| | powder_recovery.parquet | 592 | Powder collection data | -| | xrd_measurements.parquet | 592 | XRD measurement metadata | -| | dosing_sessions.parquet | 588 | Dosing session parameters | -| | sample_finalization.parquet | 592 | Labeling & storage | -| **Details (1:N)** | experiment_elements.parquet | 2,666 | Elements per experiment | -| | powder_doses.parquet | 2,245 | Individual dose events | -| | workflow_tasks.parquet | 3,074 | Task execution history | -| | temperature_logs.parquet\* | 501K | Temperature readings | -| | xrd_data_points.parquet\* | 4.7M | Raw XRD patterns | -| **Analysis** | xrd_refinements.parquet | 576 | All XRD results (success + failures) | -| | xrd_phases.parquet | 139 | Identified phases (successes only) | - -\* _Optional: Use `--fast` flag to skip these large arrays_ - -**Why multiple files?** Efficient loading - only read what you need. Most queries load <100 KB instead of 27 MB. - -📖 **See [PARQUET_STRUCTURE.md](./PARQUET_STRUCTURE.md) for complete details on structure, relationships, and usage examples.** - -## Files - -``` -data/pipeline/ -├── run_pipeline.py # Product-based CLI (create, run, list) -├── product_pipeline.py # Pipeline orchestration per product -├── pipeline_state.py # State tracking (Parquet-based) -├── s3_uploader.py # S3 OpenData uploader -├── archive/ # Deprecated MPContribs API files -│ ├── mpcontribs_manager.py -│ └── mpcontribs_uploader.py -├── pipeline_runs.parquet # Run history (auto-generated) -├── pipeline.log # Execution logs -└── requirements.txt # Dependencies -``` - -## Configuration - -### AWS Credentials - -Set in `.env` at project root for S3 uploads: - -``` -aws_access_key_id=your_key_here -aws_secret_access_key=your_secret_here -``` - -Or export as environment variables: - -```bash -export AWS_ACCESS_KEY_ID=your_key_here -export AWS_SECRET_ACCESS_KEY=your_secret_here -``` - -## CLI Usage - -### Product Management - -```bash -# Create a new data product (interactive) -python run_pipeline.py create - -# List all available products -python run_pipeline.py list - -# Show status of a product -python run_pipeline.py status --product reaction_genome - -# Validate product configuration -python run_pipeline.py validate --product reaction_genome - -# Regenerate Pydantic schema from config -python run_pipeline.py regenerate --product reaction_genome -``` - -### Running Pipelines - -```bash -# Dry run - preview what would be uploaded -python run_pipeline.py run --product reaction_genome - -# Live upload to S3 OpenData -python run_pipeline.py run --product reaction_genome --upload - -# Run specific pipeline stages only -python run_pipeline.py run --product reaction_genome --stages filter transform analyze -``` - -### Using the Shell Script - -```bash -# Use the convenience script (activates venv automatically) -./run_product_pipeline.sh create -./run_product_pipeline.sh list -./run_product_pipeline.sh run --product reaction_genome -./run_product_pipeline.sh run --product reaction_genome --upload -``` - -## State Tracking - -Pipeline state is stored in `pipeline_runs.parquet`: - -| Column | Description | -| ----------------- | ------------------------------------------------------ | -| `run_id` | Unique run identifier | -| `run_timestamp` | When the run occurred | -| `phase` | Pipeline phase (mongodb_sync, xrd_analysis, s3_upload) | -| `experiment_name` | Experiment processed | -| `status` | success / failed / pending | -| `dry_run` | Whether this was a dry run | -| `uploaded_to_s3` | Upload status | -| `s3_key` | S3 object key if uploaded | - -### Incremental Processing - -The pipeline automatically tracks: - -- **New experiments**: Not yet processed -- **Updated experiments**: `last_updated` changed since last run -- **Pending uploads**: Analyzed but not uploaded - -## S3 OpenData Integration - -### Parquet Upload - -Parquet files are uploaded directly to AWS S3 OpenData bucket: - -- **Bucket**: `s3://materialsproject-contribs/alab_synthesis/{product_name}/` -- **Format**: Parquet files with embedded metadata in schema -- **Access**: Available via AWS S3 API and HTTP - -### Files Uploaded - -All product parquet files except very large ones: - -- `experiments.parquet` - Main experiment data -- `experiment_elements.parquet` - Element compositions -- `powder_doses.parquet` - Dosing details -- `heating_sessions.parquet` - Heating parameters -- `xrd_refinements.parquet` - XRD analysis results -- `xrd_phases.parquet` - Identified phases -- And more... - -**Excluded** (too large): `temperature_logs.parquet`, `xrd_data_points.parquet` - -### Metadata - -Metadata is embedded in parquet schema following MPContribs conventions: - -```python -# Embedded in arrow table metadata -{ - "project.name": "product_name", - "project.type": "alab_synthesis", - "columns.heating_temperature.unit": "degC", - "columns.heating_temperature.description": "Heating temperature" -} -``` - -## Dashboard Integration - -The dashboard automatically reads pipeline status: - -```python -from parquet_data_loader import ParquetDataLoader - -loader = ParquetDataLoader() -status = loader.get_pipeline_status() -# {'total_experiments': 576, 'analyzed_experiments': 450, 'uploaded_experiments': 200, ...} -``` - -## Troubleshooting - -### "No experiments to upload" - -Make sure your product has experiments configured. Check with: - -```bash -python run_pipeline.py status --product your_product -``` - -### "AWS credentials not set" - -Check `.env` file exists at project root with: - -``` -aws_access_key_id=your_key_here -aws_secret_access_key=your_secret_here -``` - -### "Import error: boto3" - -Install the AWS client: - -```bash -pip install boto3 python-dotenv -``` - -### View Logs - -```bash -# Pipeline log -cat data/pipeline/pipeline.log - -# XRD analysis log -cat data/xrd_creation/batch_analysis.log -``` - -## Dependencies - -All dependencies are consolidated in the main requirements file: - -```bash -pip install -r data/requirements.txt -``` - -Key dependencies for pipeline: - -- `pandas>=2.0` - Data processing -- `pyarrow>=14.0` - Parquet files -- `boto3>=1.40.0` - S3 uploads -- `python-dotenv>=1.0` - Environment variables -- `pydantic>=2.0` - Data validation -- `PyYAML>=6.0` - Configuration files diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/__init__.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/__init__.py deleted file mode 100644 index 0d09c9bd4..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -""" -A-Lab Data Pipeline - Product-Based Pipeline System - -Orchestrates data flow: MongoDB → Parquet → XRD Analysis → S3 OpenData - -Usage: - python run_pipeline.py create # Create new product - python run_pipeline.py list # List products - python run_pipeline.py run --product # Run pipeline (dry run) - python run_pipeline.py run --product --upload # Live upload to S3 - python run_pipeline.py status --product # Show status -""" - -from .pipeline_state import PipelineStateManager -from .s3_uploader import S3Uploader, upload_product_to_s3 - -__all__ = ['PipelineStateManager', 'S3Uploader', 'upload_product_to_s3'] - diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/mpcontribs_setup.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/mpcontribs_setup.py deleted file mode 100644 index 9cfc2a2b3..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/mpcontribs_setup.py +++ /dev/null @@ -1,334 +0,0 @@ -#!/usr/bin/env python3 -""" -MPContribs Project Setup - -Creates/updates empty MPContribs project with metadata only. -Does NOT upload data - that goes to S3. - -This module handles: -- Creating new MPContribs projects -- Updating project metadata (title, authors, description, references) -- No column definitions -- No data upload -""" - -import os -import logging -from typing import Dict, Optional, List -from pathlib import Path - -logger = logging.getLogger(__name__) - - -class MPContribsProjectManager: - """Manage MPContribs project metadata (empty config only)""" - - def __init__(self, api_key: Optional[str] = None): - """ - Initialize MPContribs manager. - - Args: - api_key: Materials Project API key (or set MP_API_KEY env var) - """ - self.api_key = api_key or os.getenv('MP_API_KEY') or os.getenv('ALAB_MP_API_KEY') - - if not self.api_key: - raise ValueError( - "MP API key required. Set MP_API_KEY or ALAB_MP_API_KEY environment variable, " - "or pass api_key parameter" - ) - - self._mpr = None - self._client = None - - def _get_mprester(self): - """Get MPRester instance (lazy load)""" - if self._mpr is None: - try: - from mp_api.client import MPRester - self._mpr = MPRester(api_key=self.api_key) - except ImportError: - raise ImportError( - "mp-api package required. Install with: pip install mp-api>=0.45.13" - ) - return self._mpr - - def _get_client(self, project_name: str): - """Get ContribsClient instance (lazy load)""" - try: - from mpcontribs.client import Client as ContribsClient - return ContribsClient(project=project_name, apikey=self.api_key) - except ImportError: - raise ImportError( - "mpcontribs-client package required. Install with: pip install mpcontribs-client>=5.10.4" - ) - - def project_exists(self, project_name: str) -> bool: - """ - Check if project exists in MPContribs. - - Args: - project_name: Name of the project - - Returns: - True if project exists, False otherwise - """ - try: - mpr = self._get_mprester() - projects = mpr.contribs.get_projects() - return project_name in [p.get('name') for p in projects] - except Exception as e: - logger.error(f"Error checking project existence: {e}") - return False - - def create_project( - self, - name: str, - title: str, - authors: str, - description: str, - references: Optional[List[Dict[str, str]]] = None, - dry_run: bool = False - ) -> bool: - """ - Create empty MPContribs project with metadata. - - Args: - name: Project name (identifier) - title: Project title - authors: Comma-separated author names - description: Project description - references: Optional list of references [{"label": "...", "url": "..."}] - dry_run: If True, only simulate the operation - - Returns: - True if successful, False otherwise - """ - if dry_run: - logger.info(f"[DRY RUN] Would create MPContribs project: {name}") - logger.info(f" Title: {title}") - logger.info(f" Authors: {authors}") - logger.info(f" Description: {description[:100]}...") - if references: - logger.info(f" References: {len(references)} links") - return True - - try: - mpr = self._get_mprester() - - # Check if project already exists - if self.project_exists(name): - logger.info(f"Project '{name}' already exists, updating instead") - return self.update_project(name, title, authors, description, references) - - # Create project - logger.info(f"Creating MPContribs project: {name}") - mpr.contribs.create_project( - name=name, - title=title, - authors=authors, - description=description - ) - - # Add references if provided - if references: - client = self._get_client(name) - client.update_project({"references": references}) - logger.info(f" Added {len(references)} references") - - logger.info(f"✓ Successfully created project: {name}") - logger.info(f" View at: https://next-gen.materialsproject.org/contribs/projects/{name}") - return True - - except Exception as e: - logger.error(f"Error creating project: {e}") - return False - - def update_project( - self, - name: str, - title: Optional[str] = None, - authors: Optional[str] = None, - description: Optional[str] = None, - references: Optional[List[Dict[str, str]]] = None, - dry_run: bool = False - ) -> bool: - """ - Update existing MPContribs project metadata. - - Args: - name: Project name - title: New title (optional) - authors: New authors (optional) - description: New description (optional) - references: New references (optional) - dry_run: If True, only simulate the operation - - Returns: - True if successful, False otherwise - """ - if dry_run: - logger.info(f"[DRY RUN] Would update MPContribs project: {name}") - if title: - logger.info(f" Title: {title}") - if authors: - logger.info(f" Authors: {authors}") - if description: - logger.info(f" Description: {description[:100]}...") - if references: - logger.info(f" References: {len(references)} links") - return True - - try: - if not self.project_exists(name): - logger.error(f"Project '{name}' does not exist. Create it first.") - return False - - client = self._get_client(name) - - # Build update dict - updates = {} - if title: - updates['title'] = title - if authors: - updates['authors'] = authors - if description: - updates['description'] = description - if references: - updates['references'] = references - - if not updates: - logger.info("No updates to apply") - return True - - logger.info(f"Updating MPContribs project: {name}") - client.update_project(updates) - logger.info(f"✓ Successfully updated project: {name}") - return True - - except Exception as e: - logger.error(f"Error updating project: {e}") - return False - - def delete_project(self, name: str, dry_run: bool = False) -> bool: - """ - Delete MPContribs project. - - Args: - name: Project name - dry_run: If True, only simulate the operation - - Returns: - True if successful, False otherwise - """ - if dry_run: - logger.info(f"[DRY RUN] Would delete MPContribs project: {name}") - return True - - try: - if not self.project_exists(name): - logger.warning(f"Project '{name}' does not exist") - return False - - client = self._get_client(name) - logger.warning(f"Deleting MPContribs project: {name}") - client.delete_project() - logger.info(f"✓ Successfully deleted project: {name}") - return True - - except Exception as e: - logger.error(f"Error deleting project: {e}") - return False - - -def setup_mpcontribs_project( - product_config: Dict, - dry_run: bool = False, - api_key: Optional[str] = None -) -> bool: - """ - Setup MPContribs project from product configuration. - - Args: - product_config: Product configuration dict - dry_run: If True, simulate the operation - api_key: MP API key (optional) - - Returns: - True if successful, False otherwise - """ - manager = MPContribsProjectManager(api_key=api_key) - - # Extract metadata from product config - name = product_config.get('name') - metadata = product_config.get('metadata', {}) - - title = metadata.get('title') or f"A-Lab {name.replace('_', ' ').title()}" - authors = metadata.get('authors') or "A-Lab Team" - description = metadata.get('description') or f"Automated synthesis data for {name}" - - # Build references list - references = [] - if metadata.get('doi'): - references.append({ - 'label': 'doi', - 'url': f"https://doi.org/{metadata['doi']}" - }) - - # Add any additional references from config - if metadata.get('references'): - for ref in metadata['references']: - if isinstance(ref, dict) and 'label' in ref and 'url' in ref: - references.append(ref) - - # Always add S3 data location reference - from config_loader import get_config - config = get_config() - s3_url = f"s3://{config.s3_bucket}/{config.s3_prefix}/{name}/" - references.append({ - 'label': 'data', - 'url': s3_url.replace('s3://', 'https://s3.amazonaws.com/') - }) - - # Create/update project - return manager.create_project( - name=name, - title=title, - authors=authors, - description=description, - references=references if references else None, - dry_run=dry_run - ) - - -if __name__ == '__main__': - # Test script - import sys - - if len(sys.argv) < 2: - print("Usage: python mpcontribs_setup.py [--dry-run]") - sys.exit(1) - - project_name = sys.argv[1] - dry_run = '--dry-run' in sys.argv - - manager = MPContribsProjectManager() - - # Test project creation - success = manager.create_project( - name=project_name, - title=f"Test Project {project_name}", - authors="Test Author", - description="Test description for MPContribs project", - references=[ - {"label": "github", "url": "https://github.com/example/repo"} - ], - dry_run=dry_run - ) - - if success: - print(f"✓ Project setup successful") - else: - print(f"✗ Project setup failed") - sys.exit(1) - diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/pipeline_state.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/pipeline_state.py deleted file mode 100755 index 928a5c199..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/pipeline_state.py +++ /dev/null @@ -1,292 +0,0 @@ -#!/usr/bin/env python3 -""" -Pipeline State Manager - Tracks all pipeline runs in Parquet format - -The pipeline_runs.parquet file serves as an audit log and enables: -- Incremental processing (only process new/updated experiments) -- Upload tracking (know what's been uploaded to MPContribs) -- Run history and debugging -""" - -import uuid -from datetime import datetime -from pathlib import Path -from typing import List, Optional, Dict -import pandas as pd -import numpy as np - -PIPELINE_DIR = Path(__file__).parent -STATE_FILE = PIPELINE_DIR / "pipeline_runs.parquet" -PARQUET_DATA_DIR = PIPELINE_DIR.parent / "parquet" - - -class PipelineStateManager: - """Manage pipeline state in Parquet format""" - - def __init__(self, state_file: Path = STATE_FILE): - self.state_file = state_file - self._runs_df = None - - @property - def runs(self) -> pd.DataFrame: - """Lazy load runs dataframe""" - if self._runs_df is None: - if self.state_file.exists(): - self._runs_df = pd.read_parquet(self.state_file) - else: - self._runs_df = pd.DataFrame(columns=[ - 'run_id', 'run_timestamp', 'run_type', 'phase', - 'experiment_name', 'experiment_last_updated', - 'status', 'error_message', 'duration_seconds', - 'dry_run', 'uploaded_to_mpcontribs', 'mpcontribs_contribution_id' - ]) - return self._runs_df - - def save(self): - """Save state to parquet""" - self.runs.to_parquet(self.state_file, index=False) - - def get_experiments_to_process(self) -> List[str]: - """ - Get experiments that need processing based on: - 1. New experiments (not in state) - 2. Updated experiments (last_updated changed) - """ - experiments_file = PARQUET_DATA_DIR / "experiments.parquet" - if not experiments_file.exists(): - return [] - - experiments_df = pd.read_parquet(experiments_file) - - # Get latest successful run per experiment - if len(self.runs) == 0: - return experiments_df['name'].tolist() - - successful_runs = self.runs[ - (self.runs['phase'] == 'xrd_analysis') & - (self.runs['status'] == 'success') - ] - - if len(successful_runs) == 0: - return experiments_df['name'].tolist() - - latest_per_exp = successful_runs.groupby('experiment_name').agg({ - 'experiment_last_updated': 'max' - }).reset_index() - - # Merge with current experiments - merged = experiments_df.merge( - latest_per_exp, - left_on='name', - right_on='experiment_name', - how='left', - suffixes=('', '_processed') - ) - - # Find new or updated experiments - needs_processing = merged[ - (merged['experiment_last_updated'].isna()) | - (pd.to_datetime(merged['last_updated']) > pd.to_datetime(merged['experiment_last_updated'])) - ] - - return needs_processing['name'].tolist() - - def get_experiments_to_upload(self) -> List[str]: - """Get experiments that have been analyzed but not uploaded to MPContribs""" - # Get all experiments with successful XRD refinements - refinements_file = PARQUET_DATA_DIR / "xrd_refinements.parquet" - if not refinements_file.exists(): - return [] - - refinements_df = pd.read_parquet(refinements_file) - analyzed = set(refinements_df[refinements_df['success'] == True]['experiment_name'].unique()) - - if len(self.runs) == 0: - return list(analyzed) - - # Get uploaded experiments (non-dry-run) - uploaded = set(self.runs[ - (self.runs['phase'] == 'mpcontribs_upload') & - (self.runs['status'] == 'success') & - (self.runs['dry_run'] == False) - ]['experiment_name'].unique()) - - return list(analyzed - uploaded) - - def record_run( - self, - run_type: str, - phases: List[str], - experiments: Optional[List[str]] = None, - dry_run: bool = True - ) -> str: - """ - Record a pipeline run (batch record for all phases/experiments) - - Returns: - run_id for tracking - """ - run_id = str(uuid.uuid4())[:8] - timestamp = datetime.now() - - if experiments is None: - experiments = ['_batch_'] # Placeholder for batch operations - - new_rows = [] - for exp in experiments: - for phase in phases: - new_rows.append({ - 'run_id': run_id, - 'run_timestamp': timestamp, - 'run_type': run_type, - 'phase': phase, - 'experiment_name': exp, - 'experiment_last_updated': timestamp, - 'status': 'pending', - 'error_message': None, - 'duration_seconds': 0.0, - 'dry_run': dry_run, - 'uploaded_to_mpcontribs': False, - 'mpcontribs_contribution_id': None - }) - - if new_rows: - new_df = pd.DataFrame(new_rows) - if len(self.runs) == 0: - self._runs_df = new_df - else: - self._runs_df = pd.concat([self.runs, new_df], ignore_index=True) - self.save() - - return run_id - - def update_experiment_status( - self, - run_id: str, - experiment_name: str, - phase: str, - status: str, - duration_seconds: float = 0.0, - error_message: Optional[str] = None, - mpcontribs_id: Optional[str] = None - ): - """Update status for a specific experiment/phase in a run""" - mask = ( - (self.runs['run_id'] == run_id) & - (self.runs['experiment_name'] == experiment_name) & - (self.runs['phase'] == phase) - ) - - if not mask.any(): - # Add new row if doesn't exist - new_row = { - 'run_id': run_id, - 'run_timestamp': datetime.now(), - 'run_type': 'manual', - 'phase': phase, - 'experiment_name': experiment_name, - 'experiment_last_updated': datetime.now(), - 'status': status, - 'error_message': error_message, - 'duration_seconds': duration_seconds, - 'dry_run': False, - 'uploaded_to_mpcontribs': mpcontribs_id is not None, - 'mpcontribs_contribution_id': mpcontribs_id - } - self._runs_df = pd.concat([self.runs, pd.DataFrame([new_row])], ignore_index=True) - else: - self._runs_df.loc[mask, 'status'] = status - self._runs_df.loc[mask, 'duration_seconds'] = duration_seconds - - if error_message: - self._runs_df.loc[mask, 'error_message'] = error_message - - if mpcontribs_id: - self._runs_df.loc[mask, 'mpcontribs_contribution_id'] = mpcontribs_id - self._runs_df.loc[mask, 'uploaded_to_mpcontribs'] = True - - self.save() - - def get_run_summary(self, run_id: Optional[str] = None) -> Dict: - """Get summary statistics for a run (or latest run if not specified)""" - if len(self.runs) == 0: - return {'error': 'No runs found'} - - if run_id is None: - run_id = self.runs['run_id'].iloc[-1] - - run_data = self.runs[self.runs['run_id'] == run_id] - - if len(run_data) == 0: - return {'error': f'Run {run_id} not found'} - - return { - 'run_id': run_id, - 'timestamp': str(run_data['run_timestamp'].iloc[0]), - 'total_experiments': run_data['experiment_name'].nunique(), - 'phases': run_data['phase'].unique().tolist(), - 'success_count': len(run_data[run_data['status'] == 'success']), - 'failed_count': len(run_data[run_data['status'] == 'failed']), - 'pending_count': len(run_data[run_data['status'] == 'pending']), - 'dry_run': bool(run_data['dry_run'].iloc[0]), - 'total_duration': float(run_data['duration_seconds'].sum()) - } - - def get_last_run_timestamp(self) -> Optional[datetime]: - """Get timestamp of last successful run""" - if len(self.runs) == 0: - return None - - successful = self.runs[self.runs['status'] == 'success'] - if len(successful) == 0: - return None - - return successful['run_timestamp'].max() - - def get_status_summary(self) -> Dict: - """Get overall pipeline status summary""" - experiments_file = PARQUET_DATA_DIR / "experiments.parquet" - refinements_file = PARQUET_DATA_DIR / "xrd_refinements.parquet" - - total_experiments = 0 - analyzed_experiments = 0 - - if experiments_file.exists(): - total_experiments = len(pd.read_parquet(experiments_file)) - - if refinements_file.exists(): - ref_df = pd.read_parquet(refinements_file) - analyzed_experiments = len(ref_df[ref_df['success'] == True]) - - uploaded_experiments = 0 - if len(self.runs) > 0: - uploaded_experiments = len(self.runs[ - (self.runs['phase'] == 'mpcontribs_upload') & - (self.runs['status'] == 'success') & - (self.runs['dry_run'] == False) - ]['experiment_name'].unique()) - - return { - 'total_experiments': total_experiments, - 'analyzed_experiments': analyzed_experiments, - 'uploaded_experiments': uploaded_experiments, - 'pending_analysis': len(self.get_experiments_to_process()), - 'pending_upload': len(self.get_experiments_to_upload()), - 'last_run': str(self.get_last_run_timestamp()) if self.get_last_run_timestamp() else None, - 'total_runs': self.runs['run_id'].nunique() if len(self.runs) > 0 else 0 - } - - -if __name__ == '__main__': - mgr = PipelineStateManager() - - print("Pipeline State Summary") - print("=" * 50) - - status = mgr.get_status_summary() - for key, value in status.items(): - print(f" {key}: {value}") - - print("\nExperiments to process:", len(mgr.get_experiments_to_process())) - print("Experiments to upload:", len(mgr.get_experiments_to_upload())) - diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/product_pipeline.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/product_pipeline.py deleted file mode 100644 index 6c44b3058..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/product_pipeline.py +++ /dev/null @@ -1,691 +0,0 @@ -#!/usr/bin/env python3 -""" -Product Pipeline Runner - -Orchestrates the full pipeline for a data product: -1. Load product configuration -2. Filter experiments based on criteria -3. Transform MongoDB to Parquet (consolidated schema) -4. Run configured analyses -5. Validate with Pydantic -6. Upload parquet files directly to S3 OpenData - -Output Files: -- experiments.parquet: Main table with ALL 1:1 data (~45 columns) -- experiment_elements.parquet: Elements per experiment (1:N) -- powder_doses.parquet: Individual powder doses (1:N) -- temperature_logs.parquet: Temperature readings (1:N, optional) -- xrd_data_points.parquet: Raw XRD patterns (1:N, optional) -- workflow_tasks.parquet: Task execution history (1:N, optional) -""" - -import argparse -import sys -import logging -from pathlib import Path -from datetime import datetime -from typing import Dict, List, Optional -import pandas as pd -import shutil - -# Add parent directories to path -sys.path.insert(0, str(Path(__file__).parent.parent)) -sys.path.insert(0, str(Path(__file__).parent.parent / "products")) -sys.path.insert(0, str(Path(__file__).parent.parent / "analyses")) -sys.path.insert(0, str(Path(__file__).parent.parent / "config")) - -from base_product import ProductConfig, ProductManager -from base_analyzer import AnalysisPluginManager -from mongodb_to_parquet import MongoToParquetTransformer -from pipeline_state import PipelineStateManager -from s3_uploader import S3Uploader, upload_product_to_s3 -from mpcontribs_setup import setup_mpcontribs_project -from config_loader import get_config - -# Add tools directory for diagram generation -sys.path.insert(0, str(Path(__file__).parent.parent / "tools")) - -logging.basicConfig( - level=logging.INFO, - format='%(asctime)s - %(levelname)s - %(message)s' -) -logger = logging.getLogger(__name__) - -PARQUET_DATA_DIR = Path("data/parquet") - - -class ProductPipeline: - """Manages pipeline execution for a data product""" - - def __init__(self, product_name: str): - self.product_name = product_name - self.product_dir = Path(f"data/products/{product_name}") - self.config: Optional[ProductConfig] = None - self.state_manager = PipelineStateManager() - - def load_config(self) -> bool: - """Load product configuration""" - config_file = self.product_dir / "config.yaml" - - if not config_file.exists(): - logger.error(f"Config file not found: {config_file}") - return False - - try: - self.config = ProductConfig.load(config_file) - logger.info(f"Loaded config for product: {self.product_name}") - return True - except Exception as e: - logger.error(f"Failed to load config: {e}") - return False - - def filter_experiments(self) -> pd.DataFrame: - """ - Filter experiments from MongoDB based on configuration - - Returns: - DataFrame with filtered experiments - """ - from pymongo import MongoClient - - # Connect to MongoDB (using config loader) - config = get_config() - client = MongoClient(config.mongo_uri) - db = client[config.mongo_db] - collection = db[config.mongo_collection] - - # Build query from filter - query = self.config.experiment_filter.to_mongo_query() - - # Get matching experiments - experiments = list(collection.find(query, { - '_id': 1, - 'name': 1, - 'metadata.target': 1, - 'last_updated': 1, - 'status': 1 - })) - - client.close() - - # Convert to DataFrame - df = pd.DataFrame(experiments) - df['experiment_id'] = df['_id'].astype(str) - - # Extract target formula from metadata - df['target_formula'] = df.get('metadata', pd.Series()).apply( - lambda x: x.get('target', '') if isinstance(x, dict) else '' - ) - - # Extract experiment_type and experiment_subgroup from name - def extract_type_and_subgroup(name): - parts = name.split('_') - exp_type = parts[0] if parts else name - exp_subgroup = '_'.join(parts[:2]) if len(parts) >= 2 else exp_type - return pd.Series({'experiment_type': exp_type, 'experiment_subgroup': exp_subgroup}) - - df[['experiment_type', 'experiment_subgroup']] = df['name'].apply(extract_type_and_subgroup) - - logger.info(f"Found {len(df)} experiments matching filter") - - # Save list to product directory - experiment_list_file = self.product_dir / "experiments.txt" - df['name'].to_csv(experiment_list_file, index=False, header=False) - logger.info(f"Saved experiment list to {experiment_list_file}") - - return df - - def transform_to_parquet(self, experiments_df: pd.DataFrame) -> bool: - """ - Transform filtered experiments to consolidated Parquet format - - Uses the same transformation as update_data.sh for consistency. - """ - experiment_names = experiments_df['name'].tolist() - product_parquet_dir = self.product_dir / "parquet_data" - - # Check if product parquet already exists with correct experiments - metadata_file = product_parquet_dir / "metadata.json" - if metadata_file.exists(): - import json - with open(metadata_file) as f: - metadata = json.load(f) - existing_count = metadata.get('datasets', {}).get('experiments', {}).get('rows', 0) - if existing_count == len(experiment_names): - logger.info(f"✓ Product parquet already exists ({len(experiment_names)} experiments)") - logger.info(" Skipping transformation - using existing data") - return True - - # Create product-specific filtered parquet - logger.info(f"Creating filtered parquet for {len(experiment_names)} experiments...") - - experiment_filter = {'experiment_names': experiment_names} - - try: - transformer = MongoToParquetTransformer(output_dir=product_parquet_dir) - transformer.transform_all(experiment_filter=experiment_filter) - transformer.close() - - logger.info(f"✓ Created product-specific parquet ({len(experiment_names)} experiments)") - logger.info(f" Location: {product_parquet_dir}") - return True - - except Exception as e: - logger.error(f"Parquet transformation failed: {e}") - return False - - def run_analyses(self, experiments_df: pd.DataFrame) -> pd.DataFrame: - """ - Run configured analyses on parquet data - - Args: - experiments_df: Experiments to analyze - - Returns: - DataFrame with analysis results (if any) - """ - parquet_dir = self.product_dir / "parquet_data" - output_dir = self.product_dir / "analysis_results" - - # Ensure output directory exists - output_dir.mkdir(parents=True, exist_ok=True) - - # If no analyses configured, just return - if not self.config.analyses: - logger.info("No analyses configured") - return pd.DataFrame() - - # Run analyses using plugin manager - plugin_manager = AnalysisPluginManager() - - results_df = plugin_manager.run_analyses( - self.config.analyses, - experiments_df, - parquet_dir, - output_dir - ) - - logger.info(f"✓ Completed {len(self.config.analyses)} analyses") - - # Save analysis results (only if there are results) - if len(results_df) > 0 and len(self.config.analyses) > 0: - results_file = output_dir / "analysis_results.parquet" - results_df.to_parquet(results_file, index=False) - logger.info(f"✓ Saved analysis results: {len(results_df)} rows") - - return results_df - - def load_experiments_data(self) -> pd.DataFrame: - """ - Load experiment data from consolidated parquet file - - Returns: - DataFrame with all experiment data from experiments.parquet - """ - experiments_file = self.product_dir / "parquet_data" / "experiments.parquet" - - if not experiments_file.exists(): - logger.error(f"Experiments file not found: {experiments_file}") - return pd.DataFrame() - - return pd.read_parquet(experiments_file) - - def validate_data(self, data_df: pd.DataFrame) -> bool: - """ - Validate data against Pydantic schema - - Args: - data_df: Data to validate - - Returns: - True if all records valid - """ - # Load Pydantic schema - schema_file = self.product_dir / "schema.py" - - if not schema_file.exists(): - logger.warning("No Pydantic schema found, skipping validation") - return True - - # Import the schema module - import importlib.util - spec = importlib.util.spec_from_file_location("schema", schema_file) - schema_module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(schema_module) - - # Get the model class (assumes it follows naming convention) - model_name = f"{self.product_name.title().replace('_', '')}Experiment" - - if not hasattr(schema_module, model_name): - logger.warning(f"Schema class {model_name} not found") - return True - - model_class = getattr(schema_module, model_name) - - valid_count = 0 - error_count = 0 - - for _, row in data_df.iterrows(): - try: - # Convert row to dict - record = row.to_dict() - - # Validate against model - model_instance = model_class(**record) - valid_count += 1 - - except Exception as e: - error_count += 1 - if error_count <= 10: # Only show first 10 errors - logger.warning(f"Validation error for {row.get('name', 'unknown')}: {e}") - - logger.info(f"Validation: {valid_count} valid, {error_count} errors") - - return error_count == 0 - - def get_parquet_files(self) -> List[Path]: - """ - Get list of parquet files to upload - - Returns: - List of parquet file paths - """ - parquet_dir = self.product_dir / "parquet_data" - - if not parquet_dir.exists(): - return [] - - # Get all parquet files - parquet_files = list(parquet_dir.glob("*.parquet")) - - return parquet_files - - def setup_mpcontribs_project(self, dry_run: bool = False) -> bool: - """ - Setup empty MPContribs project with metadata only. - Does NOT upload data - that goes to S3. - - Args: - dry_run: If True, only simulate the operation - - Returns: - True if successful - """ - try: - # Convert config to dict if it's a Pydantic model - if hasattr(self.config, 'dict'): - config_dict = self.config.dict() - else: - config_dict = self.config - - return setup_mpcontribs_project( - product_config=config_dict, - dry_run=dry_run - ) - - except Exception as e: - logger.error(f"Error setting up MPContribs project: {e}") - import traceback - traceback.print_exc() - return False - - def upload_to_s3(self, dry_run: bool = True) -> bool: - """ - Upload parquet files directly to S3 OpenData. - - This is the new upload method that bypasses MPContribs library - and uploads parquet files directly to AWS S3. - - Args: - dry_run: If True, only simulate upload - - Returns: - True if successful - """ - parquet_dir = self.product_dir / "parquet_data" - - if not parquet_dir.exists(): - logger.error(f"Parquet directory not found: {parquet_dir}") - return False - - # Get parquet files (exclude very large ones by default) - parquet_files = list(parquet_dir.glob("*.parquet")) - - if not parquet_files: - logger.warning("No parquet files to upload") - return True - - # Show excluded fields info - if hasattr(self, 'schema_manager') and self.schema_manager: - logger.info("\n[Excluded from upload (embargoed):]") - for table in self.schema_manager.get_table_names(): - excluded = self.schema_manager.get_excluded_fields(table) - if excluded: - logger.info(f" {table}: {', '.join(excluded)}") - - # Show files that would be uploaded - large_files = {'temperature_logs.parquet', 'xrd_data_points.parquet'} - upload_files = [f for f in parquet_files if f.name not in large_files] - skipped_files = [f for f in parquet_files if f.name in large_files] - - total_size = sum(f.stat().st_size for f in upload_files) - - logger.info(f"\n{'[DRY RUN] ' if dry_run else ''}Files to upload ({len(upload_files)} files, {total_size / (1024*1024):.2f} MB):") - for pf in sorted(upload_files, key=lambda f: f.name): - size_kb = pf.stat().st_size / 1024 - logger.info(f" • {pf.name} ({size_kb:.1f} KB)") - - if skipped_files: - logger.info(f"\nSkipped (too large for OpenData):") - for pf in skipped_files: - size_mb = pf.stat().st_size / (1024 * 1024) - logger.info(f" • {pf.name} ({size_mb:.1f} MB)") - - if dry_run: - logger.info(f"\n[DRY RUN] Would upload {len(upload_files)} files to s3://materialsproject-contribs/alab_synthesis/{self.product_name}/") - return True - - # Perform actual upload - try: - uploaded = upload_product_to_s3( - product_name=self.product_name, - product_dir=self.product_dir, - metadata={'project': {'name': self.product_name, 'type': 'alab_synthesis'}}, - exclude_large=True, - dry_run=False - ) - - logger.info(f"✓ Uploaded {len(uploaded)} files to S3") - - # Record in state - if hasattr(self, 'state_manager') and self.state_manager: - self.state_manager.record_run( - run_type='s3_upload', - phases=['s3_upload'], - experiments=[self.product_name], - dry_run=False - ) - - return True - - except Exception as e: - logger.error(f"Failed to upload to S3: {e}") - return False - - def load_schemas(self): - """Load Pydantic schemas (schemas are now Python-first, not generated)""" - try: - from schema_manager import SchemaManager - self.schema_manager = SchemaManager() - - # List loaded schemas - schemas = self.schema_manager.get_table_names() - logger.info(f"✓ Loaded {len(schemas)} Pydantic schemas") - - # Show excluded fields - for table in schemas: - excluded = self.schema_manager.get_excluded_fields(table) - if excluded: - logger.info(f" {table}: {len(excluded)} fields excluded from upload ({', '.join(excluded)})") - - except Exception as e: - logger.warning(f"Could not load schemas: {e}") - self.schema_manager = None - - def run(self, - stages: List[str] = None, - dry_run: bool = True) -> bool: - """ - Run the complete pipeline - - Args: - stages: List of stages to run ['filter', 'transform', 'analyze', 'validate', 'diagram', 'mpcontribs', 'upload'] - dry_run: If True, simulate upload without actually uploading - - Returns: - True if successful - """ - if stages is None: - stages = ['filter', 'transform', 'analyze', 'validate', 'diagram', 'upload'] - - logger.info("=" * 60) - logger.info(f"Product Pipeline: {self.product_name}") - logger.info(f"Stages: {', '.join(stages)}") - logger.info(f"Mode: {'DRY RUN' if dry_run else 'LIVE'}") - if not dry_run and 'upload' in stages: - logger.info(f"Upload: MPContribs setup + S3 upload") - logger.info("=" * 60) - - # Load config - if not self.load_config(): - return False - - # Load Pydantic schemas (Python-first, no generation needed) - self.load_schemas() - - experiments_df = None - data_df = None - - # Stage 1: Filter experiments - if 'filter' in stages: - logger.info("\n[Stage 1/6] Filtering experiments...") - experiments_df = self.filter_experiments() - - if experiments_df.empty: - logger.error("No experiments match filter criteria") - return False - else: - # Load existing experiment list - exp_list_file = self.product_dir / "experiments.txt" - if exp_list_file.exists(): - exp_names = pd.read_csv(exp_list_file, header=None)[0].tolist() - experiments_df = pd.DataFrame({'name': exp_names}) - - # Stage 2: Transform to Parquet - if 'transform' in stages and experiments_df is not None: - logger.info("\n[Stage 2/6] Transforming to Parquet...") - if not self.transform_to_parquet(experiments_df): - return False - - # Validate schema after transformation - logger.info("\n[Stage 2.5/6] Validating schema...") - self._validate_schema() - - # Load data from consolidated experiments.parquet - data_df = self.load_experiments_data() - - # Stage 3: Run analyses - if 'analyze' in stages and experiments_df is not None: - logger.info("\n[Stage 3/6] Running analyses...") - analysis_results = self.run_analyses(experiments_df) - - # Merge analysis results with experiment data if any - if len(analysis_results) > 0: - # Merge on experiment name - data_df = data_df.merge( - analysis_results, - left_on='name', - right_on='experiment_name', - how='left' - ) - - # Stage 4: Validate - if 'validate' in stages and data_df is not None and len(data_df) > 0: - logger.info("\n[Stage 4/6] Validating with Pydantic...") - if not self.validate_data(data_df): - logger.warning("Validation errors found (continuing anyway)") - - # Stage 5: Generate Diagram - if 'diagram' in stages: - logger.info("\n[Stage 5/7] Generating schema diagram...") - self.generate_diagram() - - # Stage 6: Setup MPContribs (always when uploading for real) - if 'upload' in stages and not dry_run: - logger.info("\n[Stage 6/7] Setting up MPContribs project...") - if not self.setup_mpcontribs_project(dry_run=dry_run): - logger.warning("MPContribs setup failed (continuing anyway)") - - # Stage 7: Upload to S3 - if 'upload' in stages: - if dry_run: - logger.info("\n[Stage 7/7] Uploading to S3 OpenData (DRY RUN - no actual upload)...") - else: - logger.info("\n[Stage 7/7] Uploading to S3 OpenData...") - - if not self.upload_to_s3(dry_run=dry_run): - return False - - logger.info("\n" + "=" * 60) - logger.info("✓ Pipeline complete!") - logger.info("=" * 60) - - return True - - def _validate_schema(self): - """ - Validate parquet data using Pydantic schemas. - - Uses SchemaManager to load Pydantic schemas (Python-first, no YAML). - """ - from schema_validator import SchemaValidator - - parquet_dir = self.product_dir / "parquet_data" - - # Use pre-loaded schema manager or create new one - if not hasattr(self, 'schema_manager') or self.schema_manager is None: - from schema_manager import SchemaManager - self.schema_manager = SchemaManager() - - table_names = self.schema_manager.get_table_names() - - logger.info(f"✓ Validating {len(table_names)} tables using Pydantic schemas") - logger.info(f" Tables: {', '.join(table_names)}") - - try: - validator = SchemaValidator(self.schema_manager) - is_valid, errors, warnings = validator.validate_parquet_data(parquet_dir) - - if warnings: - from rich.console import Console - console = Console() - console.print("\n[yellow]⚠ Schema Validation Warnings:[/yellow]") - for warning in warnings[:20]: # Show first 20 - console.print(f" • {warning}") - if len(warnings) > 20: - console.print(f" • ... and {len(warnings) - 20} more") - - if errors: - from rich.console import Console - console = Console() - console.print("\n[red]✗ Schema Validation Errors:[/red]") - for error in errors: - console.print(f" • {error}") - - if is_valid: - if warnings: - from rich.console import Console - console = Console() - console.print("\n[green]✓ Schema validation passed with warnings[/green]") - else: - logger.info("✓ Schema validation passed") - else: - logger.error("✗ Schema validation failed") - - except Exception as e: - logger.warning(f"Schema validation skipped: {e}") - - def generate_diagram(self) -> bool: - """ - Generate schema diagram for this product's parquet data. - - Creates a markdown file with: - - Table overview with row counts - - Column details for each table - - Relationship diagram (Mermaid ERD) - - Output: {product_dir}/SCHEMA_DIAGRAM.md - """ - parquet_dir = self.product_dir / "parquet_data" - output_file = self.product_dir / "SCHEMA_DIAGRAM.md" - - if not parquet_dir.exists(): - logger.warning("No parquet data found, skipping diagram generation") - return False - - try: - from generate_diagram import ParquetSchemaAnalyzer, DiagramGenerator - - logger.info(f"Generating schema diagram for {self.product_name}...") - - # Analyze schema - analyzer = ParquetSchemaAnalyzer(parquet_dir) - analysis = analyzer.analyze() - - # Generate diagram - generator = DiagramGenerator(analysis) - - # Build output content - output_lines = [] - - # Header - output_lines.append(f"# {self.product_name} Schema Diagram") - output_lines.append("") - output_lines.append(f"**Auto-generated by A-Lab Pipeline**") - output_lines.append(f"") - output_lines.append(f"- **Product**: {self.product_name}") - output_lines.append(f"- **Tables**: {analysis['table_count']}") - output_lines.append(f"- **Total Rows**: {analysis['total_rows']:,}") - output_lines.append(f"- **Generated**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") - output_lines.append("") - output_lines.append("---") - output_lines.append("") - - # Add summary - output_lines.append(generator.generate_summary()) - - # Add Mermaid ERD - output_lines.append("\n---\n") - output_lines.append("## Entity Relationship Diagram\n") - output_lines.append(generator.generate_mermaid()) - - # Write to file - with open(output_file, 'w') as f: - f.write("\n".join(output_lines)) - - logger.info(f"✓ Generated schema diagram: {output_file}") - return True - - except ImportError: - logger.warning("Diagram generator not available (missing generate_diagram.py)") - return False - except Exception as e: - logger.error(f"Failed to generate diagram: {e}") - return False - - -def main(): - """CLI for running product pipeline""" - parser = argparse.ArgumentParser(description='Run product pipeline') - parser.add_argument('product_name', help='Name of the product') - parser.add_argument('--stages', nargs='+', - choices=['filter', 'transform', 'analyze', 'validate', 'diagram', 'upload'], - default=['filter', 'transform', 'analyze', 'validate', 'diagram', 'upload'], - help='Stages to run') - parser.add_argument('--live', action='store_true', - help='Run in live mode (actually upload to MPContribs)') - parser.add_argument('--dry-run', action='store_true', default=True, - help='Run in dry-run mode (default)') - - args = parser.parse_args() - - dry_run = not args.live - - pipeline = ProductPipeline(args.product_name) - success = pipeline.run(stages=args.stages, dry_run=dry_run) - - sys.exit(0 if success else 1) - - -if __name__ == '__main__': - main() diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/run_pipeline.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/run_pipeline.py deleted file mode 100755 index 18bcb5aef..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/run_pipeline.py +++ /dev/null @@ -1,429 +0,0 @@ -#!/usr/bin/env python3 -""" -A-Lab Pipeline Runner - Product-Based Pipeline - -Main CLI for managing data products and running pipelines. - -Commands: - create - Create a new data product interactively - list - List available products - run - Run pipeline for a product - status - Show product pipeline status - validate - Validate product configuration - -Usage: - python run_pipeline.py create - python run_pipeline.py run --product reaction_genome - python run_pipeline.py run --product reaction_genome --upload -""" - -import argparse -import sys -import subprocess -from pathlib import Path -from datetime import datetime -import logging -import warnings -from rich.console import Console -from rich.table import Table -from rich.panel import Panel -import pandas as pd - -# Suppress Pydantic config warnings -warnings.filterwarnings('ignore', message='Valid config keys have changed in V2') -warnings.filterwarnings('ignore', category=UserWarning, module='pydantic') - -# Add parent directories to path -sys.path.insert(0, str(Path(__file__).parent.parent)) -sys.path.insert(0, str(Path(__file__).parent.parent / "products")) - -from base_product import ProductManager, ProductConfig -from product_pipeline import ProductPipeline -from pipeline_state import PipelineStateManager - -logging.basicConfig( - level=logging.INFO, - format='%(asctime)s - %(levelname)s - %(message)s' -) -logger = logging.getLogger(__name__) - -console = Console() - - -class PipelineCLI: - """Main CLI for pipeline operations""" - - def __init__(self): - self.product_manager = ProductManager(Path("data/products")) - self.state_manager = PipelineStateManager() - - def cmd_create(self, args): - """Create a new data product interactively""" - console.print("\n[bold cyan]Creating New Data Product[/bold cyan]\n") - - try: - config = self.product_manager.create_product_interactive() - - console.print(f"\n[green]✓ Product '{config.name}' created successfully![/green]") - console.print(f"\nConfiguration saved to: data/products/{config.name}/config.yaml") - console.print(f"Pydantic schema: data/products/{config.name}/schema.py") - - # Automatically run dry run - console.print(f"\n[bold cyan]Running pipeline dry run...[/bold cyan]\n") - - from product_pipeline import ProductPipeline - pipeline = ProductPipeline(config.name) - success = pipeline.run(dry_run=True) - - if not success: - console.print(f"\n[yellow]⚠ Dry run encountered errors[/yellow]") - console.print(f"\n[dim]Fix issues and run manually:[/dim]") - console.print(f" ./run_product_pipeline.sh run --product {config.name}") - return 1 - - console.print(f"\n[green]✓ Dry run completed successfully![/green]") - - # Ask if they want to upload for real (MPContribs + S3) - import inquirer - upload = inquirer.confirm( - message=f"Upload {config.name}? (MPContribs setup + S3 upload)", - default=False - ) - - if upload: - console.print(f"\n[bold cyan]Uploading to MPContribs...[/bold cyan]\n") - pipeline = ProductPipeline(config.name) - success = pipeline.run(dry_run=False) - - if success: - console.print(f"\n[green]✓ Successfully uploaded to MPContribs![/green]") - else: - console.print(f"\n[red]✗ Upload failed[/red]") - return 1 - else: - console.print(f"\n[dim]To upload later, run:[/dim]") - console.print(f" ./run_product_pipeline.sh run --product {config.name} --upload") - - except KeyboardInterrupt: - console.print("\n[yellow]Product creation cancelled[/yellow]") - except Exception as e: - console.print(f"\n[red]Error creating product: {e}[/red]") - import traceback - console.print(traceback.format_exc()) - return 1 - - return 0 - - def cmd_list(self, args): - """List available data products""" - products = self.product_manager.list_products() - - if not products: - console.print("[yellow]No data products found[/yellow]") - console.print("\nCreate one with: python run_pipeline.py create") - return 0 - - # Create table - table = Table(title="Available Data Products") - table.add_column("Product", style="cyan") - table.add_column("Experiments", justify="right") - table.add_column("Analyses", justify="right") - table.add_column("Last Run", style="dim") - table.add_column("Status") - - for product_name in products: - try: - # Load config - config = self.product_manager.get_product_config(product_name) - - # Get experiment count - exp_file = Path(f"data/products/{product_name}/experiments.txt") - exp_count = 0 - if exp_file.exists(): - exp_count = len(exp_file.read_text().strip().split('\n')) - - # Get enabled analyses - enabled_analyses = sum(1 for a in config.analyses if a.get('enabled', True)) - - # Get last run from state - last_run = "Never" - status = "Not run" - - # Check for recent runs in state - if len(self.state_manager.runs) > 0: - product_runs = self.state_manager.runs[ - self.state_manager.runs['experiment_name'].str.contains(product_name, na=False) - ] - if len(product_runs) > 0: - last_timestamp = product_runs['run_timestamp'].max() - last_run = pd.to_datetime(last_timestamp).strftime('%Y-%m-%d %H:%M') - - # Get status - last_status = product_runs[ - product_runs['run_timestamp'] == last_timestamp - ]['status'].iloc[0] - - if last_status == 'success': - status = "[green]✓ Success[/green]" - elif last_status == 'failed': - status = "[red]✗ Failed[/red]" - else: - status = "[yellow]⟳ Running[/yellow]" - - table.add_row( - product_name, - str(exp_count), - str(enabled_analyses), - last_run, - status - ) - - except Exception as e: - table.add_row(product_name, "?", "?", "?", f"[red]Error[/red]") - - console.print(table) - - # Show additional info - console.print(f"\n[dim]Products directory: data/products/[/dim]") - console.print(f"[dim]Run a product: python run_pipeline.py run --product [/dim]") - - return 0 - - def cmd_run(self, args): - """Run pipeline for a data product""" - if not args.product: - console.print("[red]Error: --product is required[/red]") - return 1 - - # Check if product exists - products = self.product_manager.list_products() - if args.product not in products: - console.print(f"[red]Product '{args.product}' not found[/red]") - console.print(f"\nAvailable products: {', '.join(products)}") - return 1 - - # Run pipeline - pipeline = ProductPipeline(args.product) - - dry_run = not args.upload - - success = pipeline.run( - stages=args.stages, - dry_run=dry_run - ) - - return 0 if success else 1 - - def cmd_status(self, args): - """Show status for a product or all products""" - if args.product: - products = [args.product] - else: - products = self.product_manager.list_products() - - if not products: - console.print("[yellow]No products found[/yellow]") - return 0 - - for product_name in products: - console.print(Panel(f"[bold]{product_name}[/bold]", expand=False)) - - try: - # Load config - config = self.product_manager.get_product_config(product_name) - product_dir = Path(f"data/products/{product_name}") - - # Check what exists - checks = { - 'Configuration': (product_dir / 'config.yaml').exists(), - 'Pydantic Schema': (product_dir / 'schema.py').exists(), - 'Experiment List': (product_dir / 'experiments.txt').exists(), - 'Parquet Data': (product_dir / 'parquet_data').exists(), - 'Analysis Results': (product_dir / 'analysis_results').exists(), - } - - for item, exists in checks.items(): - status = "[green]✓[/green]" if exists else "[red]✗[/red]" - console.print(f" {status} {item}") - - # Show experiment count - exp_file = product_dir / 'experiments.txt' - if exp_file.exists(): - exp_count = len(exp_file.read_text().strip().split('\n')) - console.print(f"\n Experiments: {exp_count}") - - # Show filter info - if config.experiment_filter.types: - console.print(f" Types: {', '.join(config.experiment_filter.types)}") - - # Show analyses - enabled = [a['name'] for a in config.analyses if a.get('enabled', True)] - if enabled: - console.print(f" Analyses: {', '.join(enabled)}") - - console.print() - - except Exception as e: - console.print(f" [red]Error: {e}[/red]\n") - - return 0 - - def cmd_validate(self, args): - """Validate product configuration""" - if not args.product: - console.print("[red]Error: --product is required[/red]") - return 1 - - try: - # Load config - config = self.product_manager.get_product_config(args.product) - - console.print(f"\n[bold]Validating {args.product}[/bold]\n") - - # Validate configuration structure - issues = [] - warnings = [] - - # Check required fields - if not config.experiment_filter.types and not config.experiment_filter.experiment_names: - issues.append("No experiment filter defined (types or experiment_names required)") - - if not config.data_schema: - warnings.append("No schema fields defined") - - # Check analyses exist - from analyses.base_analyzer import AnalysisPluginManager - plugin_manager = AnalysisPluginManager() - available = plugin_manager.list_analyzers() - - for analysis in config.analyses: - if analysis['name'] not in available: - warnings.append(f"Analysis '{analysis['name']}' not found in available analyzers") - - # Check Pydantic schema - schema_file = Path(f"data/products/{args.product}/schema.py") - if schema_file.exists(): - console.print(" [green]✓[/green] Pydantic schema exists") - - # Try to import it - try: - import importlib.util - spec = importlib.util.spec_from_file_location("schema", schema_file) - schema_module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(schema_module) - console.print(" [green]✓[/green] Pydantic schema is valid Python") - except Exception as e: - issues.append(f"Pydantic schema error: {e}") - else: - warnings.append("No Pydantic schema generated") - - # Report results - if issues: - console.print("\n[red]Issues found:[/red]") - for issue in issues: - console.print(f" • {issue}") - - if warnings: - console.print("\n[yellow]Warnings:[/yellow]") - for warning in warnings: - console.print(f" • {warning}") - - if not issues and not warnings: - console.print("[green]✓ Configuration is valid![/green]") - - return 1 if issues else 0 - - except Exception as e: - console.print(f"[red]Error validating product: {e}[/red]") - return 1 - - def cmd_regenerate(self, args): - """Show schema info (schemas are now Python-first, no regeneration needed)""" - console.print(f"\n[bold]Pydantic Schemas (Python-first)[/bold]\n") - console.print("[cyan]Schemas are now defined directly in Python at data/products/schema/[/cyan]") - console.print("[cyan]No regeneration needed - just edit the .py files directly.[/cyan]\n") - - try: - from schema_manager import SchemaManager - schema_manager = SchemaManager() - - console.print("[bold]Available Schemas:[/bold]") - for table_name in schema_manager.get_table_names(): - schema = schema_manager.get_schema(table_name) - excluded = schema_manager.get_excluded_fields(table_name) - - console.print(f"\n [green]{table_name}[/green] ({schema.__name__})") - console.print(f" Fields: {len(schema.model_fields)}") - if excluded: - console.print(f" [yellow]Excluded from upload: {', '.join(excluded)}[/yellow]") - - return 0 - - except Exception as e: - console.print(f"[red]Error loading schemas: {e}[/red]") - return 1 - - -def main(): - parser = argparse.ArgumentParser( - description='A-Lab Pipeline Runner - Product-Based Pipeline' - ) - - subparsers = parser.add_subparsers(dest='command', help='Commands') - - # Create command - create_parser = subparsers.add_parser('create', help='Create new data product') - - # List command - list_parser = subparsers.add_parser('list', help='List available products') - - # Run command - run_parser = subparsers.add_parser('run', help='Run pipeline for a product') - run_parser.add_argument('--product', '-p', required=True, help='Product name') - run_parser.add_argument('--stages', '-s', nargs='+', - choices=['filter', 'transform', 'analyze', 'validate', 'diagram', 'upload'], - help='Stages to run (default: all)') - run_parser.add_argument('--upload', action='store_true', - help='Upload for real: MPContribs setup + S3 upload (default is dry run)') - - # Status command - status_parser = subparsers.add_parser('status', help='Show product status') - status_parser.add_argument('--product', '-p', help='Product name (or all)') - - # Validate command - val_parser = subparsers.add_parser('validate', help='Validate product config') - val_parser.add_argument('--product', '-p', required=True, help='Product name') - - # Regenerate command - regen_parser = subparsers.add_parser('regenerate', help='Regenerate Pydantic schema from config') - regen_parser.add_argument('--product', '-p', required=True, help='Product name') - - args = parser.parse_args() - - if not args.command: - # Default to 'create' if no command specified - console.print("[cyan]No command specified, starting product creation...[/cyan]\n") - args.command = 'create' - - # Execute command - cli = PipelineCLI() - - if args.command == 'create': - return cli.cmd_create(args) - elif args.command == 'list': - return cli.cmd_list(args) - elif args.command == 'run': - return cli.cmd_run(args) - elif args.command == 'status': - return cli.cmd_status(args) - elif args.command == 'validate': - return cli.cmd_validate(args) - elif args.command == 'regenerate': - return cli.cmd_regenerate(args) - else: - parser.print_help() - return 1 - - -if __name__ == '__main__': - sys.exit(main()) diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/s3_uploader.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/s3_uploader.py deleted file mode 100644 index 6fa2c042b..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/s3_uploader.py +++ /dev/null @@ -1,420 +0,0 @@ -#!/usr/bin/env python3 -""" -S3 Uploader for MPContribs OpenData - -Direct upload of parquet files to AWS S3 for MPContribs. -Based on MPContribs documentation for large data uploads. - -Requires AWS credentials from MP staff: -- aws_access_key_id -- aws_secret_access_key - -Set in .env file: - aws_access_key_id=your_key_here - aws_secret_access_key=your_secret_here -""" - -import json -import logging -import os -import sys -from pathlib import Path -from typing import Optional, List, Dict, Any - -import pandas as pd -import pyarrow as pa -import pyarrow.parquet as pq - -# Import config loader -sys.path.insert(0, str(Path(__file__).parent.parent / "config")) -from config_loader import get_config - -logger = logging.getLogger(__name__) - -# Load configuration (env vars > yaml > defaults) -config = get_config() - - -def prepare_metadata_from_config(config: Dict[str, Any]) -> Dict[str, Dict[str, Any]]: - """ - Prepare metadata dict from product configuration for embedding in parquet. - - Args: - config: Product configuration dictionary (from config.yaml or ProductConfig) - - Returns: - Dict formatted for parquet schema metadata embedding - - Example: - >>> config = {'name': 'my_product', 'metadata': {'title': 'My Dataset'}} - >>> metadata = prepare_metadata_from_config(config) - >>> # Use with S3Uploader.upload_parquet(metadata=metadata) - """ - metadata = {} - - # Project info - project_info = { - 'name': config.get('name', 'alab_synthesis'), - 'type': 'alab_synthesis' - } - - # Add metadata if available - if 'metadata' in config: - meta = config['metadata'] - if 'title' in meta: - project_info['title'] = meta['title'] - if 'description' in meta: - project_info['description'] = meta['description'] - if 'authors' in meta: - project_info['authors'] = meta['authors'] - - metadata['project'] = project_info - - # Add column metadata if schema is available - if 'data_schema' in config: - columns = {} - for field_name, field_def in config['data_schema'].items(): - col_meta = {} - - # Add description - if isinstance(field_def, dict): - if 'description' in field_def: - col_meta['description'] = field_def['description'] - if 'unit' in field_def: - col_meta['unit'] = field_def['unit'] - elif hasattr(field_def, 'description'): - col_meta['description'] = field_def.description - if hasattr(field_def, 'unit'): - col_meta['unit'] = field_def.unit - - if col_meta: - columns[field_name] = col_meta - - if columns: - metadata['columns'] = columns - - return metadata - - -class S3Uploader: - """Upload parquet files to AWS S3 OpenData for MPContribs.""" - - def __init__(self, project_name: str, api_key_id: str = None, api_key_secret: str = None): - """ - Initialize S3 uploader. - - Args: - project_name: MPContribs project name (e.g., 'alab_synthesis') - api_key_id: AWS access key ID (or from env AWS_ACCESS_KEY_ID) - api_key_secret: AWS secret access key (or from env AWS_SECRET_ACCESS_KEY) - """ - self.project_name = project_name - - # Load credentials from env or args - try: - from dotenv import load_dotenv - env_file = Path(__file__).parent.parent.parent / ".env" - if env_file.exists(): - load_dotenv(env_file) - except ImportError: - pass - - self.api_key_id = api_key_id or os.environ.get('aws_access_key_id') or os.environ.get('AWS_ACCESS_KEY_ID') - self.api_key_secret = api_key_secret or os.environ.get('aws_secret_access_key') or os.environ.get('AWS_SECRET_ACCESS_KEY') - - self._client = None - - @property - def client(self): - """Lazy load S3 client.""" - if self._client is None: - if not self.api_key_id or not self.api_key_secret: - raise ValueError( - "AWS credentials not set. Add to .env file:\n" - " aws_access_key_id=your_key\n" - " aws_secret_access_key=your_secret" - ) - - try: - import boto3 - self._client = boto3.client( - "s3", - aws_access_key_id=self.api_key_id, - aws_secret_access_key=self.api_key_secret, - ) - logger.info("Connected to AWS S3") - except ImportError: - raise ImportError("boto3 not installed. Run: pip install boto3") - - return self._client - - def upload_parquet( - self, - local_path: Path, - key: str = None, - metadata: Dict = None, - dry_run: bool = True - ) -> str: - """ - Upload a parquet file to S3. - - Args: - local_path: Path to local parquet file - key: S3 key (default: {project_name}/{filename}) - metadata: Optional metadata to embed in parquet schema - dry_run: If True, don't actually upload - - Returns: - S3 URL of uploaded file - """ - local_path = Path(local_path) - key = key or f"{self.project_name}/{local_path.name}" - - s3_url = f"s3://{config.s3_bucket}/{key}" - - if dry_run: - size_mb = local_path.stat().st_size / (1024 * 1024) - logger.info(f"[DRY RUN] Would upload: {local_path.name} ({size_mb:.2f} MB) -> {s3_url}") - return s3_url - - # If metadata provided, embed it in the parquet file - if metadata: - table = pq.read_table(local_path) - new_metadata = { - **table.schema.metadata, - **{ - f"{key}.{sub_key}": json.dumps(v) - for key, vals in metadata.items() - for sub_key, v in vals.items() - }, - } - table = table.replace_schema_metadata(metadata=new_metadata) - - # Write to temp file with metadata - temp_path = local_path.with_suffix('.tmp.parquet') - pq.write_table(table, temp_path) - upload_path = temp_path - else: - upload_path = local_path - - try: - with open(upload_path, "rb") as f: - self.client.upload_fileobj(f, Bucket=config.s3_bucket, Key=key) - - logger.info(f"✓ Uploaded: {local_path.name} -> {s3_url}") - return s3_url - - finally: - # Clean up temp file - if metadata and temp_path.exists(): - temp_path.unlink() - - def upload_all_parquet( - self, - parquet_dir: Path, - metadata: Dict = None, - exclude_large: bool = True, - dry_run: bool = True - ) -> Dict[str, str]: - """ - Upload all parquet files from a directory. - - Args: - parquet_dir: Directory containing parquet files - metadata: Optional metadata to embed in all files - exclude_large: If True, skip temperature_logs and xrd_data_points (very large) - dry_run: If True, don't actually upload - - Returns: - Dict of filename -> S3 URL - """ - parquet_dir = Path(parquet_dir) - - if not parquet_dir.exists(): - logger.error(f"Directory not found: {parquet_dir}") - return {} - - # Find all parquet files - parquet_files = list(parquet_dir.glob("*.parquet")) - - if not parquet_files: - logger.warning(f"No parquet files found in {parquet_dir}") - return {} - - # Optionally exclude large files - large_files = {'temperature_logs.parquet', 'xrd_data_points.parquet'} - if exclude_large: - parquet_files = [f for f in parquet_files if f.name not in large_files] - logger.info(f"Excluding large files: {large_files}") - - # Sort by size (smallest first) - parquet_files.sort(key=lambda f: f.stat().st_size) - - # Calculate total size - total_size = sum(f.stat().st_size for f in parquet_files) - logger.info(f"{'[DRY RUN] ' if dry_run else ''}Uploading {len(parquet_files)} files ({total_size / (1024*1024):.2f} MB total)") - - # Upload each file - uploaded = {} - for pf in parquet_files: - s3_url = self.upload_parquet(pf, metadata=metadata, dry_run=dry_run) - uploaded[pf.name] = s3_url - - return uploaded - - def delete_file(self, key: str, dry_run: bool = True) -> bool: - """ - Delete a file from S3. - - Args: - key: S3 key (e.g., 'project_name/file.parquet') - dry_run: If True, don't actually delete - - Returns: - True if successful - """ - if dry_run: - logger.info(f"[DRY RUN] Would delete: s3://{config.s3_bucket}/{key}") - return True - - try: - self.client.delete_object(Bucket=config.s3_bucket, Key=key) - logger.info(f"✓ Deleted: s3://{config.s3_bucket}/{key}") - return True - except Exception as e: - logger.error(f"Failed to delete {key}: {e}") - return False - - def list_files(self, details: bool = False) -> List[str] | List[Dict]: - """ - List all files in the project's S3 prefix. - - Args: - details: If True, return detailed info (size, last modified, etc.) - - Returns: - List of file keys (strings) or list of dicts with details - """ - try: - response = self.client.list_objects_v2( - Bucket=config.s3_bucket, - Prefix=f"{self.project_name}/" - ) - - if details: - files = [] - for obj in response.get('Contents', []): - files.append({ - 'key': obj['Key'], - 'size_bytes': obj['Size'], - 'size_mb': obj['Size'] / (1024 * 1024), - 'last_modified': obj['LastModified'], - 'etag': obj['ETag'] - }) - return files - else: - files = [] - for obj in response.get('Contents', []): - files.append(obj['Key']) - return files - - except Exception as e: - logger.error(f"Failed to list files: {e}") - return [] - - def get_upload_url(self, key: str) -> str: - """ - Get public HTTPS URL for accessing uploaded file. - - Args: - key: S3 key (e.g., 'alab_synthesis/product/file.parquet') - - Returns: - HTTPS URL for accessing the file - """ - return f"https://{config.s3_bucket}.s3.amazonaws.com/{key}" - - -def upload_product_to_s3( - product_name: str, - product_dir: Path, - metadata: Dict = None, - exclude_large: bool = True, - dry_run: bool = True, - auto_metadata: bool = True -) -> Dict[str, str]: - """ - Upload a product's parquet files to S3. - - Args: - product_name: Name of the product (used as S3 prefix) - product_dir: Path to product directory - metadata: Optional metadata to embed (if None and auto_metadata=True, reads from config.yaml) - exclude_large: If True, skip very large files (temperature_logs, xrd_data_points) - dry_run: If True, don't actually upload - auto_metadata: If True and metadata is None, try to load from config.yaml - - Returns: - Dict of filename -> S3 URL - - Example: - >>> from pathlib import Path - >>> uploaded = upload_product_to_s3( - ... product_name='reaction_genome', - ... product_dir=Path('data/products/reaction_genome'), - ... dry_run=False - ... ) - """ - uploader = S3Uploader(project_name=f"alab_synthesis/{product_name}") - parquet_dir = product_dir / "parquet_data" - - # Auto-load metadata from config.yaml if requested - if metadata is None and auto_metadata: - config_file = product_dir / "config.yaml" - if config_file.exists(): - try: - import yaml - with open(config_file, 'r') as f: - config = yaml.safe_load(f) - metadata = prepare_metadata_from_config(config) - logger.info(f"Loaded metadata from {config_file.name}") - except Exception as e: - logger.warning(f"Could not load metadata from config: {e}") - - return uploader.upload_all_parquet( - parquet_dir=parquet_dir, - metadata=metadata, - exclude_large=exclude_large, - dry_run=dry_run - ) - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser(description="Upload parquet files to MPContribs S3") - parser.add_argument("--project", "-p", required=True, help="MPContribs project name") - parser.add_argument("--dir", "-d", required=True, help="Directory containing parquet files") - parser.add_argument("--upload", action="store_true", help="Actually upload (default is dry run)") - parser.add_argument("--include-large", action="store_true", help="Include temperature_logs and xrd_data_points") - - args = parser.parse_args() - - dry_run = not args.upload - - if dry_run: - print("=" * 60) - print("DRY RUN MODE - No files will be uploaded") - print("Use --upload to actually upload to S3") - print("=" * 60) - print() - - uploader = S3Uploader(project_name=args.project) - uploaded = uploader.upload_all_parquet( - parquet_dir=Path(args.dir), - exclude_large=not args.include_large, - dry_run=dry_run - ) - - print(f"\n{'Would upload' if dry_run else 'Uploaded'} {len(uploaded)} files") - diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/test_integrated_pipeline.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/test_integrated_pipeline.py deleted file mode 100755 index bffc8cc0a..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/pipeline/test_integrated_pipeline.py +++ /dev/null @@ -1,920 +0,0 @@ -#!/usr/bin/env python3 -""" -Comprehensive Integration Tests for A-Lab Pipeline - -Tests the complete pipeline including: -- Auto-discovery (schemas, analyses) -- Filter configurations -- Hook system -- Edge cases and error handling -- End-to-end workflow - -Usage: - python test_integrated_pipeline.py # Run all tests - python test_integrated_pipeline.py --verbose # Verbose output - pytest test_integrated_pipeline.py # Use pytest -""" - -import sys -import os -import tempfile -import shutil -from pathlib import Path -from datetime import datetime -import json -import traceback - -# Add paths -sys.path.insert(0, str(Path("data/products"))) -sys.path.insert(0, str(Path("data/pipeline"))) -sys.path.insert(0, str(Path("data/analyses"))) -sys.path.insert(0, str(Path("data"))) -sys.path.insert(0, str(Path("data/tools"))) - -# Test results tracking -test_results = [] - - -def log_test(test_name: str, passed: bool, message: str = ""): - """Track test results""" - test_results.append({ - 'test': test_name, - 'passed': passed, - 'message': message - }) - status = "✓" if passed else "✗" - print(f" {status} {test_name}") - if message and not passed: - print(f" {message}") - - -def test_header(title: str): - """Print test section header""" - print("\n" + "=" * 70) - print(f" {title}") - print("=" * 70) - - -# ============================================================================= -# Auto-Discovery Tests -# ============================================================================= - -def test_schema_auto_discovery(): - """Test that schemas are auto-discovered from schema directory""" - test_header("TEST 1: Schema Auto-Discovery") - - try: - from schema_manager import SchemaManager - - manager = SchemaManager() - discovered = manager.get_table_names() - - # Should discover at least these core schemas - expected_schemas = ['experiments', 'experiment_elements', 'powder_doses'] - - all_found = all(s in discovered for s in expected_schemas) - log_test( - "Core schemas discovered", - all_found, - f"Expected {expected_schemas}, found {discovered}" - ) - - # Test that schemas have model_fields - for schema_name in discovered: - schema_class = manager.get_schema(schema_name) - has_fields = hasattr(schema_class, 'model_fields') and len(schema_class.model_fields) > 0 - log_test( - f"Schema '{schema_name}' has fields", - has_fields, - f"Fields: {len(schema_class.model_fields) if has_fields else 0}" - ) - - # Test excluded fields detection - experiments_schema = manager.get_schema('experiments') - if experiments_schema: - excluded = manager.get_excluded_fields('experiments') - log_test( - "Excluded fields detected", - True, - f"Found {len(excluded)} excluded fields: {excluded}" - ) - - return True - - except Exception as e: - log_test("Schema auto-discovery", False, str(e)) - traceback.print_exc() - return False - - -def test_analysis_auto_discovery(): - """Test that analyses are auto-discovered from analyses directory""" - test_header("TEST 2: Analysis Auto-Discovery") - - try: - from base_analyzer import AnalysisPluginManager - - manager = AnalysisPluginManager() - discovered = manager.list_analyzers() - - # Should discover at least these built-in analyses - expected_analyses = ['xrd_dara', 'powder_statistics'] - - all_found = all(a in discovered for a in expected_analyses) - log_test( - "Built-in analyses discovered", - all_found, - f"Expected {expected_analyses}, found {discovered}" - ) - - # Test analyzer metadata - for analyzer_name in discovered: - analyzer = manager.get_analyzer(analyzer_name) - - has_analyze = hasattr(analyzer, 'analyze') and callable(analyzer.analyze) - log_test(f"Analyzer '{analyzer_name}' has analyze()", has_analyze) - - has_schema = hasattr(analyzer, 'get_output_schema') and callable(analyzer.get_output_schema) - log_test(f"Analyzer '{analyzer_name}' has get_output_schema()", has_schema) - - if has_schema: - schema = analyzer.get_output_schema() - log_test( - f"Analyzer '{analyzer_name}' schema valid", - isinstance(schema, dict) and len(schema) > 0, - f"Fields: {list(schema.keys())}" - ) - - return True - - except Exception as e: - log_test("Analysis auto-discovery", False, str(e)) - traceback.print_exc() - return False - - -def test_custom_schema_hook(): - """Test adding a custom schema via hook""" - test_header("TEST 3: Custom Schema Hook") - - try: - from schema_manager import SchemaManager - from pydantic import BaseModel, Field - - # Create a temporary custom schema - custom_schema_dir = Path("data/products/schema") - custom_schema_file = custom_schema_dir / "test_custom_schema.py" - - # Write custom schema - custom_schema_code = '''"""Test custom schema""" -from pydantic import BaseModel, Field - -class TestCustomData(BaseModel, extra="forbid"): - """Test custom data schema""" - __schema_table__ = "test_custom_data" - - experiment_id: str = Field(description="Experiment ID") - custom_field: float = Field(description="Custom measurement") -''' - - with open(custom_schema_file, 'w') as f: - f.write(custom_schema_code) - - try: - # Re-discover schemas - manager = SchemaManager() - discovered = manager.get_table_names() - - custom_found = 'test_custom_data' in discovered - log_test( - "Custom schema auto-discovered", - custom_found, - f"Found in: {discovered}" - ) - - if custom_found: - custom_schema = manager.get_schema('test_custom_data') - log_test( - "Custom schema loaded correctly", - custom_schema is not None and hasattr(custom_schema, 'model_fields'), - f"Fields: {list(custom_schema.model_fields.keys()) if custom_schema else []}" - ) - - finally: - # Cleanup - if custom_schema_file.exists(): - custom_schema_file.unlink() - - return True - - except Exception as e: - log_test("Custom schema hook", False, str(e)) - traceback.print_exc() - return False - - -def test_custom_analyzer_hook(): - """Test adding a custom analyzer via hook""" - test_header("TEST 4: Custom Analyzer Hook") - - try: - from base_analyzer import AnalysisPluginManager - - # Create a temporary custom analyzer - custom_analyzer_file = Path("data/analyses/test_custom_analyzer.py") - - # Write custom analyzer - custom_analyzer_code = '''"""Test custom analyzer""" -from pathlib import Path -import pandas as pd -from base_analyzer import BaseAnalyzer - -class TestCustomAnalyzer(BaseAnalyzer): - """Test custom analysis""" - name = "test_custom" - description = "Test custom analysis for testing" - cli_flag = "--test-custom" - - def analyze(self, experiments_df: pd.DataFrame, parquet_dir: Path) -> pd.DataFrame: - results = [] - for _, exp in experiments_df.iterrows(): - results.append({ - 'experiment_name': exp.get('name', 'test'), - 'test_metric': 42.0 - }) - return pd.DataFrame(results) - - def get_output_schema(self): - return { - 'test_metric': {'type': 'float', 'required': True, 'description': 'Test metric'} - } -''' - - with open(custom_analyzer_file, 'w') as f: - f.write(custom_analyzer_code) - - try: - # Re-discover analyzers - manager = AnalysisPluginManager() - discovered = manager.list_analyzers() - - custom_found = 'test_custom' in discovered - log_test( - "Custom analyzer auto-discovered", - custom_found, - f"Found in: {discovered}" - ) - - if custom_found: - custom_analyzer = manager.get_analyzer('test_custom') - log_test( - "Custom analyzer loaded correctly", - custom_analyzer is not None and hasattr(custom_analyzer, 'analyze'), - f"Name: {custom_analyzer.name if custom_analyzer else 'N/A'}" - ) - - # Test analyzer can be instantiated - schema = custom_analyzer.get_output_schema() - log_test( - "Custom analyzer schema valid", - isinstance(schema, dict) and 'test_metric' in schema, - f"Schema: {schema}" - ) - - finally: - # Cleanup - if custom_analyzer_file.exists(): - custom_analyzer_file.unlink() - - return True - - except Exception as e: - log_test("Custom analyzer hook", False, str(e)) - traceback.print_exc() - return False - - -# ============================================================================= -# Filter Configuration Tests -# ============================================================================= - -def test_filter_configurations(): - """Test various filter configurations""" - test_header("TEST 5: Filter Configurations") - - try: - from base_product import ExperimentFilter - - # Test 1: Simple type filter - filter1 = ExperimentFilter(types=["NSC"]) - query1 = filter1.to_mongo_query() - log_test( - "Simple type filter", - 'name' in query1 and '$regex' in query1['name'], - f"Query: {query1}" - ) - - # Test 2: Multiple types - filter2 = ExperimentFilter(types=["NSC", "Na"]) - query2 = filter2.to_mongo_query() - log_test( - "Multiple type filter", - 'name' in query2, - f"Query: {query2}" - ) - - # Test 3: Status filter - filter3 = ExperimentFilter(status=["completed"]) - query3 = filter3.to_mongo_query() - log_test( - "Status filter", - 'status' in query3, - f"Query: {query3}" - ) - - # Test 4: XRD requirement - filter4 = ExperimentFilter(has_xrd=True) - query4 = filter4.to_mongo_query() - log_test( - "XRD requirement filter", - 'metadata.diffraction_results.sampleid_in_aeris' in query4, - f"Query: {query4}" - ) - - # Test 5: Combined filters - filter5 = ExperimentFilter( - types=["NSC"], - status=["completed"], - has_xrd=True - ) - query5 = filter5.to_mongo_query() - log_test( - "Combined filters", - len(query5) >= 2, - f"Query has {len(query5)} conditions" - ) - - # Test 6: Specific experiment names - filter6 = ExperimentFilter(experiment_names=["NSC_249_001", "NSC_249_002"]) - query6 = filter6.to_mongo_query() - log_test( - "Specific experiment names", - 'name' in query6 and '$in' in query6['name'], - f"Query: {query6}" - ) - - # Test 7: Date range filter - filter7 = ExperimentFilter( - date_range={ - "start": "2024-01-01", - "end": "2024-12-31" - } - ) - query7 = filter7.to_mongo_query() - log_test( - "Date range filter", - 'last_updated' in query7, - f"Query: {query7}" - ) - - # Test 8: Empty filter (should match all) - filter8 = ExperimentFilter() - query8 = filter8.to_mongo_query() - log_test( - "Empty filter (match all)", - query8 == {}, - f"Query: {query8}" - ) - - return True - - except Exception as e: - log_test("Filter configurations", False, str(e)) - traceback.print_exc() - return False - - -# ============================================================================= -# Edge Cases and Error Handling -# ============================================================================= - -def test_edge_cases(): - """Test edge cases and potential breaking scenarios""" - test_header("TEST 6: Edge Cases & Error Handling") - - try: - from base_product import ProductConfig, ExperimentFilter, ProductMetadata - - # Test 1: Invalid product name - try: - invalid_config = ProductConfig( - name="invalid name with spaces!", - experiment_filter=ExperimentFilter() - ) - log_test("Invalid product name rejected", False, "Should have raised ValueError") - except ValueError: - log_test("Invalid product name rejected", True, "Correctly rejected invalid name") - - # Test 2: Empty product name - try: - empty_config = ProductConfig( - name="", - experiment_filter=ExperimentFilter() - ) - log_test("Empty product name rejected", False, "Should have raised ValueError") - except (ValueError, Exception): - log_test("Empty product name rejected", True, "Correctly rejected empty name") - - # Test 3: Valid product with underscores - try: - valid_config = ProductConfig( - name="valid_product_123", - experiment_filter=ExperimentFilter() - ) - log_test("Valid product name accepted", True, f"Name: {valid_config.name}") - except Exception as e: - log_test("Valid product name accepted", False, str(e)) - - # Test 4: Missing required fields - try: - from schema_manager import SchemaManager - manager = SchemaManager() - exp_schema = manager.get_schema('experiments') - - if exp_schema: - # Try to create with missing required fields - try: - invalid_exp = exp_schema( - experiment_id="test", - name="test" - # Missing other required fields - ) - log_test("Missing required fields rejected", False, "Should have raised validation error") - except Exception: - log_test("Missing required fields rejected", True, "Correctly rejected") - except Exception as e: - log_test("Schema validation test", False, str(e)) - - # Test 5: Field validation (optional fields accept None) - try: - from schema_manager import SchemaManager - manager = SchemaManager() - exp_schema = manager.get_schema('experiments') - - if exp_schema: - # Test that validation works at all - try: - # Try to create with invalid status (not in Literal) - invalid_exp = exp_schema( - experiment_id="test", - name="test", - experiment_type="TEST", - target_formula="test", - last_updated=datetime.now(), - status="invalid_status", # Not in Literal["completed", "error", "active", "unknown"] - ) - log_test("Invalid enum values rejected", False, "Should have rejected invalid status") - except Exception: - log_test("Invalid enum values rejected", True, "Correctly rejected invalid status") - - # Test that valid data passes - try: - valid_exp = exp_schema( - experiment_id="test", - name="test", - experiment_type="TEST", - target_formula="test", - last_updated=datetime.now(), - status="completed", - heating_temperature=1100.0 # Valid temperature - ) - log_test("Valid data accepted", True, "Schema accepts valid data") - except Exception as e: - log_test("Valid data accepted", False, f"Schema rejected valid data: {e}") - except Exception as e: - log_test("Schema validation test", False, str(e)) - - # Test 6: Non-existent analyzer - try: - from base_analyzer import AnalysisPluginManager - manager = AnalysisPluginManager() - - non_existent = manager.get_analyzer("this_does_not_exist") - log_test( - "Non-existent analyzer handled", - non_existent is None, - "Should return None for missing analyzer" - ) - except Exception as e: - log_test("Non-existent analyzer test", False, str(e)) - - # Test 7: Non-existent schema - try: - from schema_manager import SchemaManager - manager = SchemaManager() - - non_existent = manager.get_schema("this_does_not_exist") - log_test( - "Non-existent schema handled", - non_existent is None, - "Should return None for missing schema" - ) - except Exception as e: - log_test("Non-existent schema test", False, str(e)) - - return True - - except Exception as e: - log_test("Edge cases", False, str(e)) - traceback.print_exc() - return False - - -def test_config_files(): - """Test configuration files exist and are valid""" - test_header("TEST 7: Configuration Files") - - try: - import yaml - - # Test defaults.yaml - defaults_file = Path("data/config/defaults.yaml") - if defaults_file.exists(): - with open(defaults_file) as f: - defaults = yaml.safe_load(f) - - log_test( - "defaults.yaml exists and valid", - isinstance(defaults, dict) and 'version' in defaults, - f"Version: {defaults.get('version')}" - ) - - # Check required sections - required_sections = ['mongodb', 'parquet', 'analyses', 'upload'] - for section in required_sections: - log_test( - f"defaults.yaml has '{section}' section", - section in defaults, - f"Keys: {list(defaults.keys())}" - ) - else: - log_test("defaults.yaml exists", False, "File not found") - - # Test filters.yaml - filters_file = Path("data/config/filters.yaml") - if filters_file.exists(): - with open(filters_file) as f: - filters = yaml.safe_load(f) - - log_test( - "filters.yaml exists and valid", - isinstance(filters, dict), - f"Keys: {list(filters.keys())}" - ) - - # Check presets exist - if 'presets' in filters: - presets = filters['presets'] - log_test( - "Filter presets defined", - len(presets) > 0, - f"Presets: {list(presets.keys())}" - ) - else: - log_test("filters.yaml exists", False, "File not found") - - # Test analyses.yaml - analyses_file = Path("data/config/analyses.yaml") - if analyses_file.exists(): - with open(analyses_file) as f: - analyses = yaml.safe_load(f) - - log_test( - "analyses.yaml exists and valid", - isinstance(analyses, dict) and 'analyses' in analyses, - f"Keys: {list(analyses.keys())}" - ) - - # Check built-in analyses documented - if 'analyses' in analyses: - builtin = analyses['analyses'] - expected_analyses = ['xrd_dara', 'powder_statistics'] - for analyzer in expected_analyses: - log_test( - f"Analysis '{analyzer}' documented", - analyzer in builtin, - f"Documented: {list(builtin.keys())}" - ) - else: - log_test("analyses.yaml exists", False, "File not found") - - return True - - except Exception as e: - log_test("Configuration files", False, str(e)) - traceback.print_exc() - return False - - -def test_mongodb_connection(): - """Test MongoDB connection (optional)""" - test_header("TEST 8: MongoDB Connection (Optional)") - - try: - from pymongo import MongoClient - - client = MongoClient("mongodb://localhost:27017/", serverSelectionTimeoutMS=2000) - db = client["temporary"] - collection = db["release"] - - # Try to count documents - count = collection.count_documents({}) - log_test( - "MongoDB connection successful", - count >= 0, - f"Found {count} experiments in database" - ) - - # Test aggregation pipeline - try: - pipeline = [ - {"$group": {"_id": {"$substr": ["$name", 0, 3]}, "count": {"$sum": 1}}}, - {"$sort": {"count": -1}}, - {"$limit": 5} - ] - results = list(collection.aggregate(pipeline)) - log_test( - "MongoDB aggregation works", - len(results) > 0, - f"Top types: {[r['_id'] for r in results]}" - ) - except Exception as e: - log_test("MongoDB aggregation", False, str(e)) - - client.close() - return True - - except Exception as e: - log_test("MongoDB connection", False, f"Not available: {e}") - print(" ℹ MongoDB tests skipped (server not available)") - return True # Don't fail if MongoDB isn't running - - -def test_parquet_transformer(): - """Test parquet transformation with filters""" - test_header("TEST 9: Parquet Transformer with Filters") - - try: - from mongodb_to_parquet import MongoToParquetTransformer - - # Create a temp directory for test output - with tempfile.TemporaryDirectory() as tmpdir: - output_dir = Path(tmpdir) / "test_parquet" - - # Test with filter (limit to 5 experiments for speed) - experiment_filter = { - 'status': ['completed'], - 'has_xrd': True - } - - try: - transformer = MongoToParquetTransformer(output_dir=output_dir) - # Use limit to make test fast - transformer.transform_all( - limit=5, - skip_temp_logs=True, - skip_xrd_points=True, - experiment_filter=experiment_filter - ) - transformer.close() - - # Check output files - metadata.json should always exist - metadata_file = output_dir / 'metadata.json' - log_test( - "Generated metadata.json", - metadata_file.exists(), - f"Size: {metadata_file.stat().st_size if metadata_file.exists() else 0} bytes" - ) - - # Parquet files may be empty if no data matches filter - # This is OK - just log for informational purposes - if (output_dir / 'experiments.parquet').exists(): - size = (output_dir / 'experiments.parquet').stat().st_size - log_test( - "Generated experiments.parquet", - True, - f"Size: {size} bytes (0 bytes = no matching data, OK)" - ) - else: - log_test("experiments.parquet not created", True, "No matching data (OK)") - - log_test("Parquet transformation with filters", True, "Successfully generated filtered parquet") - - except Exception as e: - log_test("Parquet transformation", False, f"MongoDB not available or error: {e}") - - return True - - except Exception as e: - log_test("Parquet transformer", False, str(e)) - print(" ℹ Transformer test skipped (MongoDB may not be available)") - return True - - -def test_diagram_generation(): - """Test schema diagram generation""" - test_header("TEST 10: Diagram Generation") - - try: - # Check if we have parquet data to diagram - parquet_dir = Path("data/parquet") - - if not parquet_dir.exists() or not list(parquet_dir.glob("*.parquet")): - log_test( - "Diagram generation", - True, - "Skipped: No parquet data available" - ) - return True - - from generate_diagram import ParquetSchemaAnalyzer, DiagramGenerator - - # Analyze schema - analyzer = ParquetSchemaAnalyzer(parquet_dir) - analysis = analyzer.analyze() - - log_test( - "Schema analysis", - 'tables' in analysis and len(analysis['tables']) > 0, - f"Found {len(analysis['tables'])} tables" - ) - - # Generate diagram - generator = DiagramGenerator(analysis) - - # Test terminal output (shouldn't error) - try: - mermaid = generator.generate_mermaid() - log_test( - "Mermaid ERD generation", - '```mermaid' in mermaid and 'erDiagram' in mermaid, - f"Generated {len(mermaid)} characters" - ) - except Exception as e: - log_test("Mermaid generation", False, str(e)) - - # Test summary generation - try: - summary = generator.generate_summary() - log_test( - "Summary generation", - '# Parquet Schema Documentation' in summary, - f"Generated {len(summary)} characters" - ) - except Exception as e: - log_test("Summary generation", False, str(e)) - - return True - - except Exception as e: - log_test("Diagram generation", False, str(e)) - traceback.print_exc() - return False - - -def test_full_integration(): - """Test full pipeline integration""" - test_header("TEST 11: Full Pipeline Integration") - - try: - # Check all required directories exist - dirs_to_check = [ - Path("data/products"), - Path("data/products/schema"), - Path("data/pipeline"), - Path("data/analyses"), - Path("data/config"), - Path("data"), - Path("data/tools") - ] - - all_exist = True - for dir_path in dirs_to_check: - exists = dir_path.exists() - log_test(f"Directory {dir_path} exists", exists) - if not exists: - all_exist = False - - # Check core files - files_to_check = [ - Path("data/products/schema_manager.py"), - Path("data/products/schema_validator.py"), - Path("data/analyses/base_analyzer.py"), - Path("data/pipeline/product_pipeline.py"), - Path("data/config/defaults.yaml"), - Path("data/config/filters.yaml"), - Path("data/config/analyses.yaml"), - Path("run_product_pipeline.sh") - ] - - for file_path in files_to_check: - exists = file_path.exists() - log_test(f"File {file_path.name} exists", exists) - if not exists: - all_exist = False - - # Check Python dependencies - try: - import pandas - import pydantic - import yaml - import rich - log_test("Core Python dependencies installed", True) - except ImportError as e: - log_test("Core Python dependencies", False, str(e)) - all_exist = False - - return all_exist - - except Exception as e: - log_test("Full integration", False, str(e)) - traceback.print_exc() - return False - - -# ============================================================================= -# Main Test Runner -# ============================================================================= - -def main(): - """Run all integration tests""" - print("\n" + "=" * 70) - print(" A-LAB PIPELINE COMPREHENSIVE INTEGRATION TESTS") - print("=" * 70) - print(f"\n Started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") - - # Run all tests - tests = [ - test_schema_auto_discovery, - test_analysis_auto_discovery, - test_custom_schema_hook, - test_custom_analyzer_hook, - test_filter_configurations, - test_edge_cases, - test_config_files, - test_mongodb_connection, - test_parquet_transformer, - test_diagram_generation, - test_full_integration - ] - - total_tests = 0 - passed_tests = 0 - - for test_func in tests: - try: - result = test_func() - # Count individual test results logged within each function - except Exception as e: - print(f"\n✗ Test suite {test_func.__name__} crashed: {e}") - traceback.print_exc() - - # Summary - print("\n" + "=" * 70) - print(" TEST SUMMARY") - print("=" * 70) - - passed = sum(1 for r in test_results if r['passed']) - failed = sum(1 for r in test_results if not r['passed']) - total = len(test_results) - - print(f"\n Total: {total} tests") - print(f" Passed: {passed} tests ({'green' if passed == total else 'yellow'})") - print(f" Failed: {failed} tests ({'red' if failed > 0 else 'green'})") - - if failed > 0: - print("\n Failed tests:") - for result in test_results: - if not result['passed']: - print(f" ✗ {result['test']}") - if result['message']: - print(f" {result['message']}") - - print("\n" + "=" * 70) - - if failed == 0: - print(" ✓ ALL TESTS PASSED!") - print("\n Next steps:") - print(" 1. Create a product: ./run_product_pipeline.sh create") - print(" 2. Run pipeline: ./run_product_pipeline.sh run --product ") - print("=" * 70) - return 0 - else: - print(" ✗ SOME TESTS FAILED") - print("\n Review failures above and fix issues.") - print("=" * 70) - return 1 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/SCHEMA_SYSTEM.md b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/SCHEMA_SYSTEM.md deleted file mode 100644 index 1b1530656..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/SCHEMA_SYSTEM.md +++ /dev/null @@ -1,202 +0,0 @@ -# A-Lab Schema System (Pydantic-First) - -This document describes the Pydantic-first schema system for A-Lab data validation and MPContribs uploads. - -## Overview - -**Schemas are defined directly in Python using Pydantic** - not generated from YAML. This approach: - -- Aligns with the team's existing `results_schema.py` -- Enables rich validation (constraints, validators, Literal types) -- Better IDE support (autocomplete, type checking) -- Matches MPContribs expectations - -## Schema Files - -All schemas are in `data/products/schema/`: - -``` -data/products/schema/ -├── __init__.py # Package exports -├── base.py # ExcludeFromUpload utility -├── experiments.py # Main consolidated table (~37 fields) -├── experiment_elements.py # Elements per experiment (1:N) -├── powder_doses.py # Individual powder doses (1:N) -├── temperature_logs.py # Temperature readings (1:N, optional) -├── xrd_data_points.py # Raw XRD patterns (1:N, optional) -├── workflow_tasks.py # Task execution history (1:N) -├── xrd_refinements.py # DARA analysis results -└── xrd_phases.py # Identified crystal phases -``` - -**One schema per parquet file** - no redundancy. -All constraints from team's `results_schema.py` are integrated into these schemas. - -## Parquet Table Schemas - -Each schema corresponds to **one parquet file**. All constraints from team's `results_schema.py` are integrated: - -| Schema | Class | Table | Relationship | Constraints | -| ------------------------ | --------------------- | --------------------------- | ------------ | ------------------------ | -| `experiments.py` | `Experiment` | experiments.parquet | Main table | Literal types, ranges | -| `experiment_elements.py` | `ExperimentElement` | experiment_elements.parquet | 1:N | - | -| `powder_doses.py` | `PowderDose` | powder_doses.parquet | 1:N | ge=0 for masses | -| `temperature_logs.py` | `TemperatureLogEntry` | temperature_logs.parquet | 1:N | - | -| `workflow_tasks.py` | `WorkflowTask` | workflow_tasks.parquet | 1:N | Literal status | -| `xrd_data_points.py` | `XRDDataPoint` | xrd_data_points.parquet | 1:N | ge=0 for point_index | -| `xrd_refinements.py` | `XRDRefinement` | xrd_refinements.parquet | 1:1 | - | -| `xrd_phases.py` | `XRDPhase` | xrd_phases.parquet | 1:N | 0 <= weight_fraction <=1 | - -### Data Flow - -MongoDB (nested) → Flattened → Parquet (normalized) → Schemas validate parquet - -The team's `results_schema.py` describes MongoDB structure. Our schemas describe the **flattened parquet** representation of that same data. - -## Field Exclusion (Embargoed Data) - -Some fields should NOT be uploaded to MPContribs until the paper is published: - -```python -from .base import ExcludeFromUpload - -class RecoverPowderResult(BaseModel, extra="forbid"): - # EXCLUDED FROM UPLOAD - embargoed until paper publication - weight_collected: float | None = ExcludeFromUpload( - description="Weight of powder collected (EMBARGOED)" - ) -``` - -**Currently excluded fields:** - -- `recovery_weight_collected_mg` - powder recovery weight -- `xrd_total_mass_dispensed_mg` - XRD mass dispensed -- `mass_per_dispensing_attempt_mg` - per-attempt XRD mass -- `first_tapping_mass_collected` - first tapping mass - -## Using Schemas - -### SchemaManager - -```python -from data.products.schema_manager import SchemaManager - -sm = SchemaManager() - -# List all tables -print(sm.get_table_names()) -# ['experiments', 'experiment_elements', 'powder_doses', ...] - -# Get schema class -Experiment = sm.get_schema('experiments') - -# Get fields for upload (excluding embargoed) -uploadable = sm.get_uploadable_fields('experiments') -# ['experiment_id', 'name', ...] (35 of 37 fields) - -excluded = sm.get_excluded_fields('experiments') -# ['recovery_weight_collected_mg', 'xrd_total_mass_dispensed_mg'] -``` - -### Validation - -```python -from data.products.schema_validator import SchemaValidator - -validator = SchemaValidator() -is_valid, errors, warnings = validator.validate_parquet_data(parquet_dir) -``` - -### Direct Import - -```python -from data.products.schema import ( - Experiment, - PowderDose, - XRDRefinement, - XRDPhase, -) - -# Validate a row -exp = Experiment( - experiment_id="abc123", - name="NSC_249", - experiment_type="NSC", - target_formula="Na3V2(PO4)3", - status="completed", - last_updated=datetime.now(), - # Heating data (flattened from MongoDB metadata.heating_results) - heating_temperature=800.0, - heating_time=240.0, - # Recovery data (flattened from MongoDB metadata.recoverpowder_results) - recovery_initial_crucible_weight_mg=1500.0, - # Note: recovery_weight_collected_mg is EXCLUDED from upload (embargoed) -) -``` - -## Upload to S3 OpenData - -The pipeline now uploads parquet files directly to S3: - -```python -from data.pipeline.s3_uploader import upload_product_to_s3 - -# Upload (dry run) -uploaded = upload_product_to_s3( - product_name="my_product", - product_dir=Path("data/products/my_product"), - dry_run=True -) - -# Actual upload -uploaded = upload_product_to_s3(..., dry_run=False) -``` - -Files are uploaded to: `s3://materialsproject-contribs/alab_synthesis/{product_name}/` - -## Modifying Schemas - -To add or modify fields, **edit the Python files directly**: - -1. Open `data/products/schema/experiments.py` -2. Add/modify fields with Pydantic syntax -3. Run pipeline - validation will use updated schema automatically - -```python -# Add a new field -new_field: str | None = Field( - default=None, - description="Description of the new field" -) - -# Mark a field as excluded from upload -sensitive_field: float | None = ExcludeFromUpload( - description="Sensitive data (EMBARGOED)" -) - -# Add constraints -position: int | None = Field( - default=None, - description="Position index", - ge=1, # >= 1 - le=16 # <= 16 -) - -# Use Literal for enum-like fields -status: Literal["completed", "error", "active", "unknown"] = Field(...) -``` - -## Parquet-Schema Alignment - -The `parquet_data_loader.py` dashboard loader is aligned with the schemas: - -| Schema Field | Parquet Column | -| ------------------------------ | ------------------------------ | -| `heating_temperature` | `heating_temperature` | -| `heating_time` | `heating_time` | -| `heating_cooling_rate` | `heating_cooling_rate` | -| `recovery_weight_collected_mg` | `recovery_weight_collected_mg` | -| `xrd_sampleid_in_aeris` | `xrd_sampleid_in_aeris` | -| `dosing_crucible_position` | `dosing_crucible_position` | - -All prefixed fields (heating*, recovery*, xrd*, dosing*, finalization\_) are consistent between schemas and parquet files. diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/base_product.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/base_product.py deleted file mode 100644 index c4f0f87e5..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/base_product.py +++ /dev/null @@ -1,771 +0,0 @@ -#!/usr/bin/env python3 -""" -Base Product Configuration System for A-Lab Data Products - -Provides a flexible framework for defining data products with: -- Experiment filtering (by type, status, date, etc.) -- Schema definition with units for MPContribs -- Pydantic validation -- Analysis pipeline configuration -- Metadata for publications -""" - -import yaml -import json -import sys -from pathlib import Path -from datetime import datetime -from typing import Dict, List, Optional, Any, Type -from pydantic import BaseModel, Field, validator, create_model -from enum import Enum -import pandas as pd -import subprocess -from rich.console import Console - -# Import config loader -sys.path.insert(0, str(Path(__file__).parent.parent / "config")) -from config_loader import get_config - -console = Console() - - -class ExperimentType(str, Enum): - """Known experiment type prefixes""" - NSC = "NSC" - Na = "Na" - PG = "PG" - MINES = "MINES" - TRI = "TRI" - - -class ExperimentStatus(str, Enum): - """Experiment workflow status""" - COMPLETED = "completed" - ERROR = "error" - ACTIVE = "active" - UNKNOWN = "unknown" - - -class AnalysisConfig(BaseModel): - """Configuration for a single analysis module""" - name: str - enabled: bool = True - config: Dict[str, Any] = {} - - -class SchemaField(BaseModel): - """Schema field definition for MPContribs""" - unit: Optional[str] = Field(None, description="Unit for numeric fields, '' for strings, None for dimensionless") - description: str = Field(..., description="Human-readable description") - type: str = Field("string", description="Data type: float, int, boolean, string") - required: bool = Field(False, description="Whether field is required") - min: Optional[float] = None - max: Optional[float] = None - - -class ExperimentFilter(BaseModel): - """Filter criteria for selecting experiments""" - types: Optional[List[ExperimentType]] = Field(None, description="Experiment type prefixes") - status: Optional[List[ExperimentStatus]] = Field(None, description="Workflow status") - has_xrd: Optional[bool] = Field(None, description="Must have XRD data") - has_sem: Optional[bool] = Field(None, description="Must have SEM data") - date_range: Optional[Dict[str, str]] = Field(None, description="Start/end dates") - experiment_names: Optional[List[str]] = Field(None, description="Specific experiment names") - - def to_mongo_query(self) -> Dict: - """Convert filter to MongoDB query""" - query = {} - - # Handle name filtering (types OR specific experiment names) - name_queries = [] - - if self.types: - # Match experiments starting with any of the prefixes - # Convert enum values to strings - type_strings = [t.value if hasattr(t, 'value') else str(t) for t in self.types] - patterns = [f"^{t}_" for t in type_strings] - name_queries.append({"name": {"$regex": "|".join(patterns)}}) - - if self.experiment_names: - # Specific experiment names - name_queries.append({"name": {"$in": self.experiment_names}}) - - # Combine name queries with $or if both exist - if len(name_queries) == 1: - query.update(name_queries[0]) - elif len(name_queries) > 1: - query["$or"] = name_queries - - if self.status: - query["status"] = {"$in": [s.value for s in self.status]} - - if self.has_xrd is not None: - if self.has_xrd: - query["metadata.diffraction_results.sampleid_in_aeris"] = {"$exists": True, "$ne": None} - else: - # Need to handle $or carefully if it already exists - xrd_or = [ - {"metadata.diffraction_results.sampleid_in_aeris": {"$exists": False}}, - {"metadata.diffraction_results.sampleid_in_aeris": None} - ] - if "$or" in query: - # Wrap both in $and - existing_or = query.pop("$or") - query["$and"] = [ - {"$or": existing_or}, - {"$or": xrd_or} - ] - else: - query["$or"] = xrd_or - - if self.has_sem is not None: - # Check for SEM data existence - if self.has_sem: - query["metadata.sem_results"] = {"$exists": True, "$ne": None} - - if self.date_range: - date_query = {} - if "start" in self.date_range: - date_query["$gte"] = datetime.fromisoformat(self.date_range["start"]) - if "end" in self.date_range: - date_query["$lte"] = datetime.fromisoformat(self.date_range["end"]) - if date_query: - query["last_updated"] = date_query - - return query - - -class ProductMetadata(BaseModel): - """Metadata for MPContribs project""" - title: Optional[str] = Field(None, description="Human-readable title") - authors: Optional[str] = Field(None, description="Author list") - description: Optional[str] = Field(None, description="Abstract/description") - references: Optional[List[Dict[str, str]]] = Field(None, description="List of {label, url}") - doi: Optional[str] = Field(None, description="DOI if published") - - -class ProductConfig(BaseModel): - """Complete configuration for a data product""" - name: str = Field(..., description="Product identifier (alphanumeric + underscore)") - version: str = Field("1.0", description="Product version") - schema_version: str = Field("1.0", description="Schema version (for migration tracking)") - - experiment_filter: ExperimentFilter - metadata: ProductMetadata = ProductMetadata() - analyses: List[AnalysisConfig] = [] - data_schema: Dict[str, SchemaField] = Field(default_factory=dict, alias='schema') - - created_at: datetime = Field(default_factory=datetime.now) - updated_at: datetime = Field(default_factory=datetime.now) - submitted_to_mpcontribs: bool = Field(False, description="Whether product has been submitted to MPContribs") - - model_config = { - 'populate_by_name': True # Allow both 'data_schema' and 'schema' to work - } - - @validator("name") - def validate_name(cls, v): - """Ensure name is valid identifier""" - if not v.replace("_", "").isalnum(): - raise ValueError("Product name must be alphanumeric with underscores only") - return v - - def save(self, path: Path): - """Save configuration to YAML file""" - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - - # Convert to dict with serializable types - config_dict = json.loads(self.json()) - - with open(path, "w") as f: - yaml.dump(config_dict, f, default_flow_style=False, sort_keys=False) - - @classmethod - def load(cls, path: Path) -> "ProductConfig": - """Load configuration from YAML file""" - with open(path, "r") as f: - data = yaml.safe_load(f) - return cls(**data) - - def to_pydantic_model(self) -> Type[BaseModel]: - """ - Generate Pydantic model from schema definition - - Returns a dynamically created Pydantic model class for validation - """ - fields = {} - validators = {} - - for field_name, field_def in self.data_schema.items(): - # Determine Python type - if field_def.type == "float": - python_type = float - elif field_def.type == "int": - python_type = int - elif field_def.type == "boolean": - python_type = bool - else: - python_type = str - - # Make optional if not required - if not field_def.required: - python_type = Optional[python_type] - - # Create field with description - fields[field_name] = (python_type, Field(description=field_def.description)) - - # Add validators for min/max - if field_def.min is not None or field_def.max is not None: - def make_validator(fname, fmin, fmax): - def validator(cls, v): - if v is not None: - if fmin is not None and v < fmin: - raise ValueError(f"{fname} must be >= {fmin}") - if fmax is not None and v > fmax: - raise ValueError(f"{fname} must be <= {fmax}") - return v - return validator - - validators[f"validate_{field_name}"] = validator(field_name)( - make_validator(field_name, field_def.min, field_def.max) - ) - - # Create dynamic model - model_name = f"{self.name.title().replace('_', '')}Experiment" - return create_model(model_name, **fields, __validators__=validators) - - def get_mpcontribs_columns(self) -> Dict[str, Optional[str]]: - """ - Get columns with units for MPContribs initialization - - Returns: - Dict of column_name -> unit (None for dimensionless, "" for strings) - """ - return { - self._to_camel_case(name): field.unit - for name, field in self.data_schema.items() - } - - @staticmethod - def _to_camel_case(snake_str: str) -> str: - """Convert snake_case to camelCase for MPContribs""" - components = snake_str.split('_') - return components[0] + ''.join(x.title() for x in components[1:]) - - -class ProductManager: - """Manager for data product lifecycle""" - - def __init__(self, products_dir: Path = None): - self.products_dir = Path(products_dir or "data/products") - self.products_dir.mkdir(parents=True, exist_ok=True) - - def create_product_interactive(self) -> ProductConfig: - """Interactive CLI for creating a new product configuration""" - import inquirer - from rich.console import Console - from rich.table import Table - from rich.tree import Tree - - console = Console() - - console.print("\n[bold cyan]Create New Data Product[/bold cyan]\n") - - # Product name - name = inquirer.text( - message="Product name (alphanumeric + underscore)", - validate=lambda _, x: x.replace("_", "").isalnum() - ) - - # Discover experiment types using MongoDB - console.print("\n[yellow]Discovering experiment types from MongoDB...[/yellow]") - hierarchy = self._discover_experiment_types() - - # Display hierarchical structure - tree = Tree("[bold cyan]Available Experiment Groups[/bold cyan]") - - for root, data in sorted(hierarchy.items()): - root_node = tree.add(f"[green]{root}[/green] ({data['count']} experiments)") - - if data['subgroups']: - for subgroup, subdata in sorted(data['subgroups'].items()): - root_node.add(f"[yellow]{subgroup}[/yellow] ({subdata['count']} experiments)") - - console.print(tree) - console.print() - - # Calculate total experiments - total_experiments = sum(data['count'] for data in hierarchy.values()) - - console.print(f"[yellow]ℹ Selection behavior:[/yellow]") - console.print(f" • Select specific groups to filter experiments") - console.print(f" • Leave empty to include ALL {total_experiments} experiments (will require confirmation)") - console.print() - - # Build selection choices with hierarchy - all_choices = [] - choice_map = {} # Maps display string to selection value - - for root in sorted(hierarchy.keys()): - data = hierarchy[root] - - # Add root level option - root_label = f"{root} (all {data['count']} experiments)" - all_choices.append(root_label) - choice_map[root_label] = {'type': 'root', 'value': root} - - # Add subgroup options - if data['subgroups']: - for subgroup in sorted(data['subgroups'].keys()): - subdata = data['subgroups'][subgroup] - subgroup_label = f" └─ {subgroup} ({subdata['count']} experiments)" - all_choices.append(subgroup_label) - choice_map[subgroup_label] = {'type': 'subgroup', 'value': subgroup} - - # Select experiment groups - selected_labels = inquirer.checkbox( - message="Select experiment groups (leave empty for ALL experiments)", - choices=all_choices - ) - - # If nothing selected, confirm they want ALL experiments - if not selected_labels: - console.print(f"\n[yellow]⚠ No groups selected - this will include ALL {total_experiments} experiments![/yellow]") - confirm_all = inquirer.confirm( - message=f"Proceed with ALL {total_experiments} experiments?", - default=False - ) - - if not confirm_all: - console.print("\n[cyan]Please select specific experiment groups...[/cyan]\n") - selected_labels = inquirer.checkbox( - message="Select experiment groups", - choices=all_choices - ) - - # If still nothing selected, abort - if not selected_labels: - console.print("\n[red]No experiments selected. Aborting product creation.[/red]") - return None - - # Process selections into types and experiment_names - selected_roots = set() - selected_subgroups = [] - selected_experiment_names = [] - - for label in selected_labels: - choice = choice_map[label] - if choice['type'] == 'root': - selected_roots.add(choice['value']) - elif choice['type'] == 'subgroup': - selected_subgroups.append(choice['value']) - # Get all experiments for this subgroup - root = choice['value'].split('_')[0] - if root in hierarchy and choice['value'] in hierarchy[root]['subgroups']: - experiments = hierarchy[root]['subgroups'][choice['value']]['experiments'] - selected_experiment_names.extend(experiments) - - # Build filter config - filter_config = {} - - # If nothing selected (user confirmed to use ALL), don't add type/name filters - # This will match all experiments in the database - if not selected_roots and not selected_experiment_names: - console.print(f"[dim]Filter: ALL experiments (no type/name filter applied)[/dim]") - # If only subgroups selected (no full roots), use experiment_names - elif selected_experiment_names and not selected_roots: - filter_config["experiment_names"] = selected_experiment_names - console.print(f"[dim]Filter: {len(selected_experiment_names)} specific experiments[/dim]") - # If roots selected, use types (and ignore subgroups under those roots) - elif selected_roots: - filter_config["types"] = list(selected_roots) - console.print(f"[dim]Filter: Types {list(selected_roots)}[/dim]") - # If both, combine: use types for roots, add experiment_names for other subgroups - elif selected_roots and selected_subgroups: - filter_config["types"] = list(selected_roots) - # Only add experiment names from subgroups whose root isn't selected - filtered_names = [ - exp for exp in selected_experiment_names - if exp.split('_')[0] not in selected_roots - ] - if filtered_names: - filter_config["experiment_names"] = filtered_names - console.print(f"[dim]Filter: Types {list(selected_roots)} + {len(filtered_names)} specific experiments[/dim]") - - # Additional filters - let user select which ones to apply - console.print("\n[cyan]Additional Filters[/cyan]") - available_filters = inquirer.checkbox( - message="Select additional filters to apply (optional)", - choices=[ - "XRD requirement", - "Status filter", - "SEM requirement", - "Date range" - ] - ) - - # Configure selected filters - if "XRD requirement" in available_filters: - xrd_options = inquirer.checkbox( - message="XRD requirement (select one)", - choices=["Must have XRD", "Must NOT have XRD"], - carousel=True - ) - if "Must have XRD" in xrd_options: - filter_config["has_xrd"] = True - elif "Must NOT have XRD" in xrd_options: - filter_config["has_xrd"] = False - - if "Status filter" in available_filters: - status = inquirer.checkbox( - message="Select experiment statuses to include", - choices=["completed", "error", "active", "unknown"], - default=["completed"] - ) - if status: - filter_config["status"] = status - - if "SEM requirement" in available_filters: - sem_options = inquirer.checkbox( - message="SEM requirement (select one)", - choices=["Must have SEM", "Must NOT have SEM"], - carousel=True - ) - if "Must have SEM" in sem_options: - filter_config["has_sem"] = True - elif "Must NOT have SEM" in sem_options: - filter_config["has_sem"] = False - - if "Date range" in available_filters: - start_date = inquirer.text( - message="Start date (YYYY-MM-DD, or blank for no start)", - default="" - ) - end_date = inquirer.text( - message="End date (YYYY-MM-DD, or blank for no end)", - default="" - ) - date_range = {} - if start_date: - date_range["start"] = start_date - if end_date: - date_range["end"] = end_date - if date_range: - filter_config["date_range"] = date_range - - # Metadata (optional) - include_metadata = inquirer.confirm( - message="Include publication metadata?", - default=False - ) - - metadata = {} - if include_metadata: - metadata["title"] = inquirer.text(message="Title") - metadata["authors"] = inquirer.text(message="Authors") - metadata["description"] = inquirer.text(message="Description") - - # Analyses - analyses = [] - available_analyses = self._discover_analyses() - - if available_analyses: - console.print("\n[cyan]Available Analyses:[/cyan]") - for analysis in available_analyses: - console.print(f" • {analysis}") - - selected_analyses = inquirer.checkbox( - message="Select analyses to run", - choices=available_analyses - ) - - analyses = [{"name": a, "enabled": True} for a in selected_analyses] - - # Show schema info (schemas are managed centrally in schema/ directory) - console.print("\n[cyan]A-Lab Pydantic Schemas[/cyan]") - from schema_manager import SchemaManager - - schema_manager = SchemaManager() - experiments_schema = schema_manager.get_main_schema() - - if experiments_schema: - # Get field info from Pydantic model - field_count = len(experiments_schema.model_fields) - excluded_fields = schema_manager.get_excluded_fields('experiments') - - console.print(f"[green]✓ Experiments schema: {field_count} fields[/green]") - if excluded_fields: - console.print(f" [yellow]⚠ {len(excluded_fields)} fields excluded from upload (embargoed)[/yellow]") - - # Show summary of fields by prefix - console.print("\n[dim]Schema includes:[/dim]") - all_fields = list(experiments_schema.model_fields.keys()) - prefixes = { - 'Core': [f for f in all_fields if f in ['experiment_id', 'name', 'experiment_type', 'target_formula', 'status']], - 'Heating': [f for f in all_fields if f.startswith('heating_')], - 'Recovery': [f for f in all_fields if f.startswith('recovery_')], - 'XRD': [f for f in all_fields if f.startswith('xrd_')], - 'Dosing': [f for f in all_fields if f.startswith('dosing_')], - 'Finalization': [f for f in all_fields if f.startswith('finalization_')] - } - - for category, fields in prefixes.items(): - if fields: - console.print(f" [yellow]{category}[/yellow]: {len(fields)} fields") - - console.print() - console.print("[dim]Note: Schemas are managed centrally in data/products/schema/[/dim]") - console.print("[dim]Edit the .py files directly to modify validation rules[/dim]") - else: - console.print(f"[yellow]⚠ Experiments schema not found[/yellow]") - - # Create config (no data_schema - managed centrally now) - config = ProductConfig( - name=name, - experiment_filter=ExperimentFilter(**filter_config), - metadata=ProductMetadata(**metadata) if metadata else ProductMetadata(), - analyses=analyses, - data_schema={} # Empty - schemas are managed centrally in schema/ directory - ) - - # Save - config_path = self.products_dir / name / "config.yaml" - config.save(config_path) - - console.print(f"\n[green]✓ Created configuration: {config_path}[/green]") - - # Show Pydantic schemas info (schemas are Python-first, not generated) - console.print(f"\n[cyan]Pydantic Schemas (from data/products/schema/):[/cyan]") - for table_name in schema_manager.get_table_names(): - schema = schema_manager.get_schema(table_name) - excluded = schema_manager.get_excluded_fields(table_name) - excluded_str = f" (excluded: {', '.join(excluded)})" if excluded else "" - console.print(f" • {table_name}: {len(schema.model_fields)} fields{excluded_str}") - - return config - - def _discover_experiment_types(self) -> Dict[str, Any]: - """ - Discover experiment types with hierarchical grouping based on underscores - - Returns: - Dict with hierarchical structure: - { - 'NSC': { - 'count': 218, - 'subgroups': { - 'NSC_249': {'count': 5, 'experiments': ['NSC_249_001', ...]}, - 'NSC_250': {'count': 3, 'experiments': [...]} - } - }, - ... - } - """ - from pymongo import MongoClient - from collections import defaultdict - - try: - # Connect to MongoDB (using config loader) - config = get_config() - client = MongoClient(config.mongo_uri, serverSelectionTimeoutMS=5000) - db = client[config.mongo_db] - collection = db[config.mongo_collection] - - # Get all experiment names - experiments = collection.find({}, {"name": 1, "_id": 0}) - - # Build hierarchical structure - hierarchy = defaultdict(lambda: { - 'count': 0, - 'subgroups': defaultdict(lambda: { - 'count': 0, - 'experiments': [] - }) - }) - - for exp in experiments: - exp_name = exp.get('name', '') - if not exp_name or '_' not in exp_name: - continue - - # Split by underscore to get all levels - parts = exp_name.split('_') - - # Root level (e.g., 'NSC', 'Na') - root = parts[0] - hierarchy[root]['count'] += 1 - - # If there are multiple parts, create subgroups - if len(parts) >= 2: - # Build all possible subgroup prefixes - # For NSC_249_001, create: NSC_249 - # For Na_123_A_001, create: Na_123, Na_123_A - for i in range(1, len(parts)): - subgroup = '_'.join(parts[:i+1]) - - # Only track subgroups that aren't the full experiment name - # (i.e., there's at least one more part after this) - if i < len(parts) - 1: - hierarchy[root]['subgroups'][subgroup]['count'] += 1 - hierarchy[root]['subgroups'][subgroup]['experiments'].append(exp_name) - - client.close() - - # Convert defaultdict to regular dict for serialization - result = {} - for root, data in hierarchy.items(): - result[root] = { - 'count': data['count'], - 'subgroups': dict(data['subgroups']) if data['subgroups'] else {} - } - - return result - - except Exception as e: - console.print(f"[yellow]Warning: Could not connect to MongoDB: {e}[/yellow]") - console.print("[yellow]Using fallback experiment types[/yellow]") - # Return empty hierarchy as fallback - return { - "NSC": {"count": 0, "subgroups": {}}, - "Na": {"count": 0, "subgroups": {}}, - "PG": {"count": 0, "subgroups": {}}, - "MINES": {"count": 0, "subgroups": {}}, - "TRI": {"count": 0, "subgroups": {}} - } - - def _discover_analyses(self) -> List[str]: - """Discover available analysis modules""" - analyses = ["xrd_dara"] # Always available - - # Check for other analysis scripts - analysis_paths = [ - Path("data/analyses/powder_statistics.py"), - Path("data/analyses/sem_clustering.py"), - Path("data/analyses/heating_profile.py") - ] - - for path in analysis_paths: - if path.exists(): - analyses.append(path.stem) - - return analyses - - def _generate_pydantic_schema(self, config: ProductConfig): - """ - Generate Pydantic validation schema file from product config - - This is auto-generated from the YAML schema and should be regenerated - whenever the schema changes. - """ - schema_path = self.products_dir / config.name / "schema.py" - - # Generate Python code for the schema - code = f'''#!/usr/bin/env python3 -""" -Auto-generated Pydantic schema for {config.name} -Generated: {datetime.now().isoformat()} - -⚠️ DO NOT EDIT THIS FILE MANUALLY ⚠️ -This file is auto-generated from config.yaml -To modify the schema, edit config.yaml and regenerate with: - python data/products/base_product.py regenerate {config.name} -""" - -from pydantic import BaseModel, Field, validator -from typing import Optional -from datetime import datetime as dt_type - - -class {config.name.title().replace("_", "")}Experiment(BaseModel): - """ - Validation schema for {config.name} experiments - - This schema validates data before upload to MPContribs. - All fields are derived from the product's data_schema in config.yaml. - """ - - # Auto-generated fields -''' - - # Add fields - for field_name, field_def in config.data_schema.items(): - python_type = { - "float": "float", - "int": "int", - "boolean": "bool", - "string": "str", - "datetime": "dt_type" - }.get(field_def.type, "str") - - if not field_def.required: - python_type = f"Optional[{python_type}]" - - default = " = None" if not field_def.required else "" - - # Add unit info to docstring if present - unit_str = f" ({field_def.unit})" if field_def.unit else "" - code += f" {field_name}: {python_type}{default} # {field_def.description}{unit_str}\n" - - # Add validators - validators_added = False - for field_name, field_def in config.data_schema.items(): - if field_def.min is not None or field_def.max is not None: - if not validators_added: - code += '\n # Validators for numeric ranges\n' - validators_added = True - - code += f''' - @validator("{field_name}") - def validate_{field_name}(cls, v): - """Validate {field_name} is within acceptable range""" - if v is not None: -''' - if field_def.min is not None: - code += f' if v < {field_def.min}:\n' - code += f' raise ValueError("{field_name} must be >= {field_def.min}")\n' - if field_def.max is not None: - code += f' if v > {field_def.max}:\n' - code += f' raise ValueError("{field_name} must be <= {field_def.max}")\n' - code += ' return v\n' - - # Add metadata - code += f''' - -# Metadata -__schema_version__ = "{config.version}" -__generated_at__ = "{datetime.now().isoformat()}" -__source_file__ = "config.yaml" -__field_count__ = {len(config.data_schema)} -''' - - # Write schema file - with open(schema_path, 'w') as f: - f.write(code) - - console.print(f"[green]✓ Generated Pydantic schema: {schema_path}[/green]") - console.print(f"[dim] {len(config.data_schema)} fields, {sum(1 for f in config.data_schema.values() if f.required)} required[/dim]") - - def list_products(self) -> List[str]: - """List all available products""" - products = [] - for path in self.products_dir.iterdir(): - if path.is_dir() and (path / "config.yaml").exists(): - products.append(path.name) - return products - - def get_product_config(self, name: str) -> ProductConfig: - """Load product configuration by name""" - config_path = self.products_dir / name / "config.yaml" - if not config_path.exists(): - raise FileNotFoundError(f"Product '{name}' not found") - return ProductConfig.load(config_path) - - -if __name__ == "__main__": - # Test the interactive product creation - manager = ProductManager() - config = manager.create_product_interactive() - print(f"\nCreated product: {config.name}") - print(f"Experiments filter: {config.experiment_filter}") - print(f"Schema fields: {list(config.data_schema.keys())}") diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/schema_manager.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/schema_manager.py deleted file mode 100644 index 857cb5321..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/schema_manager.py +++ /dev/null @@ -1,301 +0,0 @@ -#!/usr/bin/env python3 -""" -Schema Manager for A-Lab Pipeline (Pydantic-First, Auto-Discovery) - -Manages Pydantic schema classes for all parquet tables. -Schemas are defined directly in Python, NOT generated from YAML. - -EXTENSION: To add a new schema (e.g., SEM data): -1. Create a new file: data/products/schema/sem_data.py -2. Define a Pydantic class with __schema_table__ = "sem_data" -3. The schema will be auto-discovered on next pipeline run - -Schema files are located in data/products/schema/: -- experiments.py (main table, consolidated) -- experiment_elements.py -- powder_doses.py -- temperature_logs.py -- xrd_data_points.py -- workflow_tasks.py -- xrd_refinements.py -- xrd_phases.py -- [your_new_schema.py] - auto-discovered! -""" - -import importlib.util -import inspect -import logging -from pathlib import Path -from typing import Dict, List, Optional, Type -from pydantic import BaseModel - -logger = logging.getLogger(__name__) - -SCHEMA_DIR = Path(__file__).parent / "schema" - -# Files to skip during auto-discovery -SKIP_FILES = {'__init__.py', 'base.py', '__pycache__'} - -# Fallback mapping for files without __schema_table__ attribute -# Maps filename (without .py) to expected class name -LEGACY_SCHEMA_MAPPING = { - 'experiments': 'Experiment', - 'experiment_elements': 'ExperimentElement', - 'powder_doses': 'PowderDose', - 'temperature_logs': 'TemperatureLogEntry', - 'workflow_tasks': 'WorkflowTask', - 'xrd_data_points': 'XRDDataPoint', - 'xrd_refinements': 'XRDRefinement', - 'xrd_phases': 'XRDPhase', - 'heating': 'HeatingResult', - 'recovery': 'RecoverPowderResult', - 'diffraction': 'DiffractionResult', - 'powder_dosing': 'PowderDosingSampleResult', -} - - -class SchemaManager: - """ - Manages all Pydantic schema classes for the A-Lab pipeline. - - Auto-discovers schemas from the schema directory. To add a new schema: - 1. Create a .py file in data/products/schema/ - 2. Define a Pydantic BaseModel class - 3. Optionally add __schema_table__ = "table_name" to specify the table name - (otherwise derived from filename) - """ - - def __init__(self, schema_dir: Path = SCHEMA_DIR): - self.schema_dir = Path(schema_dir) - self.schemas: Dict[str, Type[BaseModel]] = {} - self._discover_schemas() - - def _discover_schemas(self): - """ - Auto-discover all Pydantic schema classes from the schema directory. - - Discovery rules: - 1. Scan all .py files in schema_dir (except SKIP_FILES) - 2. Look for classes that inherit from BaseModel - 3. Use __schema_table__ attribute if present, else derive from filename - 4. Fall back to LEGACY_SCHEMA_MAPPING for backwards compatibility - """ - if not self.schema_dir.exists(): - logger.warning(f"Schema directory not found: {self.schema_dir}") - return - - # Auto-discover all .py files - for schema_file in sorted(self.schema_dir.glob("*.py")): - if schema_file.name in SKIP_FILES: - continue - - table_name = schema_file.stem # filename without .py - - try: - # Import the module - spec = importlib.util.spec_from_file_location(table_name, schema_file) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - - # Find BaseModel subclasses in the module - schema_class = self._find_schema_class(module, table_name) - - if schema_class: - # Get table name from class attribute or filename - actual_table_name = getattr(schema_class, '__schema_table__', table_name) - self.schemas[actual_table_name] = schema_class - logger.debug(f"Discovered schema: {actual_table_name} -> {schema_class.__name__}") - - except Exception as e: - logger.error(f"Failed to load schema {schema_file}: {e}") - - logger.info(f"Auto-discovered {len(self.schemas)} schemas from {self.schema_dir}") - - def _find_schema_class(self, module, table_name: str) -> Optional[Type[BaseModel]]: - """ - Find the main schema class in a module. - - Priority: - 1. Class with __schema_table__ attribute matching table_name - 2. Class name from LEGACY_SCHEMA_MAPPING - 3. First BaseModel subclass found (excluding imports) - """ - candidates = [] - - for name, obj in inspect.getmembers(module, inspect.isclass): - # Skip imported classes (only want classes defined in this module) - if obj.__module__ != module.__name__: - continue - - # Must be a BaseModel subclass - if not issubclass(obj, BaseModel) or obj is BaseModel: - continue - - # Check for explicit table name - if hasattr(obj, '__schema_table__') and obj.__schema_table__ == table_name: - return obj - - candidates.append((name, obj)) - - # Try legacy mapping - if table_name in LEGACY_SCHEMA_MAPPING: - expected_class = LEGACY_SCHEMA_MAPPING[table_name] - for name, obj in candidates: - if name == expected_class: - return obj - - # Return first candidate if any - if candidates: - return candidates[0][1] - - return None - - def get_schema(self, table_name: str) -> Optional[Type[BaseModel]]: - """Get Pydantic schema class for a specific table""" - return self.schemas.get(table_name) - - def get_all_schemas(self) -> Dict[str, Type[BaseModel]]: - """Get all loaded schema classes""" - return self.schemas - - def get_table_names(self) -> List[str]: - """Get list of all table names""" - return sorted(self.schemas.keys()) - - def get_main_schema(self) -> Optional[Type[BaseModel]]: - """Get the main Experiment schema (consolidated table)""" - return self.schemas.get('experiments') - - def get_uploadable_fields(self, table_name: str) -> List[str]: - """ - Get list of fields that should be uploaded to MPContribs. - - Fields marked with exclude_from_upload=True are excluded. - """ - schema = self.get_schema(table_name) - if not schema: - return [] - - uploadable = [] - for field_name, field_info in schema.model_fields.items(): - extra = field_info.json_schema_extra or {} - if not extra.get("exclude_from_upload", False): - uploadable.append(field_name) - - return uploadable - - def get_excluded_fields(self, table_name: str) -> List[str]: - """ - Get list of fields that should NOT be uploaded to MPContribs. - """ - schema = self.get_schema(table_name) - if not schema: - return [] - - excluded = [] - for field_name, field_info in schema.model_fields.items(): - extra = field_info.json_schema_extra or {} - if extra.get("exclude_from_upload", False): - excluded.append(field_name) - - return excluded - - def validate_row(self, table_name: str, row_data: dict) -> tuple[bool, Optional[str]]: - """ - Validate a single row of data against its Pydantic schema. - - Args: - table_name: Name of the table - row_data: Dictionary of field values - - Returns: - (is_valid, error_message) - """ - schema = self.get_schema(table_name) - if not schema: - return True, None # No schema = skip validation - - try: - schema(**row_data) - return True, None - except Exception as e: - return False, str(e) - - def get_schema_fields(self, table_name: str) -> Dict[str, dict]: - """ - Get field information for a schema. - - Returns dict of field_name -> {type, required, description} - """ - schema = self.get_schema(table_name) - if not schema: - return {} - - fields = {} - for field_name, field_info in schema.model_fields.items(): - fields[field_name] = { - 'type': str(field_info.annotation), - 'required': field_info.is_required(), - 'description': field_info.description or '', - 'exclude_from_upload': (field_info.json_schema_extra or {}).get('exclude_from_upload', False) - } - - return fields - - def get_schema_summary(self) -> Dict[str, Dict]: - """Get summary of all schemas""" - summary = {} - for table_name, schema in self.schemas.items(): - doc = schema.__doc__ or '' - summary[table_name] = { - 'description': doc.split('\n')[0].strip() if doc else '', - 'class_name': schema.__name__, - 'num_fields': len(schema.model_fields), - 'uploadable_fields': len(self.get_uploadable_fields(table_name)), - 'excluded_fields': len(self.get_excluded_fields(table_name)) - } - return summary - - -if __name__ == '__main__': - # Test the schema manager with auto-discovery - from rich.console import Console - from rich.table import Table - - console = Console() - - manager = SchemaManager() - - console.print(f"\n[bold cyan]Schema Auto-Discovery Results[/bold cyan]") - console.print(f"Directory: {manager.schema_dir}") - console.print(f"Discovered: {len(manager.schemas)} schemas\n") - - # Create table - table = Table(title="Discovered Schemas") - table.add_column("Table Name", style="cyan") - table.add_column("Class", style="green") - table.add_column("Fields", justify="right") - table.add_column("Uploadable", justify="right") - table.add_column("Excluded", style="yellow") - - for table_name in manager.get_table_names(): - schema = manager.get_schema(table_name) - excluded = manager.get_excluded_fields(table_name) - uploadable = manager.get_uploadable_fields(table_name) - - table.add_row( - table_name, - schema.__name__, - str(len(schema.model_fields)), - str(len(uploadable)), - ", ".join(excluded) if excluded else "-" - ) - - console.print(table) - - console.print("\n[bold green]✓ Auto-discovery complete![/bold green]") - console.print("\n[dim]To add a new schema:[/dim]") - console.print(" 1. Create: data/products/schema/your_schema.py") - console.print(" 2. Define: class YourSchema(BaseModel, extra='forbid'): ...") - console.print(" 3. Optional: __schema_table__ = 'your_table_name'") - console.print(" 4. Run pipeline - schema auto-discovered!") diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/schema_validator.py b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/schema_validator.py deleted file mode 100644 index cbe0deabe..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/data/products/schema_validator.py +++ /dev/null @@ -1,217 +0,0 @@ -#!/usr/bin/env python3 -""" -Schema Validator - Validates parquet data against Pydantic schemas - -Uses the Pydantic schema classes directly for validation. -No YAML files - schemas are Python-first. -""" - -import logging -from pathlib import Path -from typing import Dict, List, Tuple, Type -import pandas as pd -from pydantic import BaseModel - -logger = logging.getLogger(__name__) - - -class SchemaValidator: - """Validates parquet data against Pydantic schemas""" - - def __init__(self, schema_manager=None): - """ - Args: - schema_manager: SchemaManager instance (will create if not provided) - """ - if schema_manager is None: - from schema_manager import SchemaManager - self.schema_manager = SchemaManager() - else: - self.schema_manager = schema_manager - - self.warnings = [] - self.errors = [] - self.validation_stats = {} - - def validate_parquet_data(self, parquet_dir: Path) -> Tuple[bool, List[str], List[str]]: - """ - Validate all parquet files against their Pydantic schemas. - - Args: - parquet_dir: Directory containing parquet files - - Returns: - (is_valid, errors, warnings) - """ - self.warnings = [] - self.errors = [] - self.validation_stats = {} - - parquet_dir = Path(parquet_dir) - - # Validate each table - for table_name in self.schema_manager.get_table_names(): - schema = self.schema_manager.get_schema(table_name) - - if not schema: - continue - - parquet_file = parquet_dir / f"{table_name}.parquet" - - if not parquet_file.exists(): - # Temperature logs and XRD data points are optional - if table_name in ['temperature_logs', 'xrd_data_points']: - logger.debug(f"Optional table not present: {table_name}.parquet") - else: - self.warnings.append(f"Table not found: {table_name}.parquet") - continue - - # Validate this table - self._validate_table(table_name, parquet_file, schema) - - is_valid = len(self.errors) == 0 - return is_valid, self.errors, self.warnings - - def _validate_table(self, table_name: str, parquet_file: Path, schema: Type[BaseModel]): - """ - Validate a parquet table against its Pydantic schema. - - Args: - table_name: Name of the table - parquet_file: Path to the parquet file - schema: Pydantic schema class - """ - try: - df = pd.read_parquet(parquet_file) - except Exception as e: - self.errors.append(f"[{table_name}] Failed to load parquet: {e}") - return - - # Check column existence - schema_fields = set(schema.model_fields.keys()) - parquet_columns = set(df.columns) - - missing_required = [] - for field_name, field_info in schema.model_fields.items(): - if field_info.is_required() and field_name not in parquet_columns: - missing_required.append(field_name) - - if missing_required: - self.errors.append( - f"[{table_name}] Missing required columns: {', '.join(missing_required)}" - ) - - # Missing optional columns are just warnings - missing_optional = schema_fields - parquet_columns - set(missing_required) - if missing_optional: - self.warnings.append( - f"[{table_name}] Missing optional columns: {', '.join(sorted(missing_optional)[:5])}" - + (f" (+{len(missing_optional)-5} more)" if len(missing_optional) > 5 else "") - ) - - # Validate a sample of rows with Pydantic - sample_size = min(100, len(df)) - sample_df = df.sample(n=sample_size, random_state=42) if len(df) > sample_size else df - - valid_count = 0 - error_count = 0 - error_samples = [] - - for idx, row in sample_df.iterrows(): - try: - row_dict = {k: v for k, v in row.to_dict().items() if pd.notna(v) or k in schema_fields} - # Handle NaN/None values - cleaned_dict = {} - for k, v in row_dict.items(): - if pd.isna(v): - cleaned_dict[k] = None - else: - cleaned_dict[k] = v - - schema(**cleaned_dict) - valid_count += 1 - except Exception as e: - error_count += 1 - if len(error_samples) < 3: - error_samples.append(f"Row {idx}: {str(e)[:150]}") - - # Store stats - self.validation_stats[table_name] = { - 'total': len(df), - 'sampled': sample_size, - 'valid': valid_count, - 'errors': error_count - } - - # Report validation errors - if error_count > 0: - self.errors.append( - f"[{table_name}] {error_count}/{sample_size} sampled rows failed validation" - ) - for sample in error_samples: - self.errors.append(f" • {sample}") - else: - logger.debug(f"[{table_name}] ✓ {valid_count} sampled rows valid") - - -def validate_product_schema(product_dir: Path, schema_manager=None) -> bool: - """ - Validate a product's parquet data against Pydantic schemas. - - Args: - product_dir: Product directory containing parquet_data/ - schema_manager: Optional SchemaManager instance - - Returns: - True if validation passes (or passes with warnings only) - """ - from rich.console import Console - console = Console() - - validator = SchemaValidator(schema_manager) - parquet_dir = product_dir / "parquet_data" - - if not parquet_dir.exists(): - console.print(f"[yellow]⚠ No parquet data found at {parquet_dir}[/yellow]") - return True - - is_valid, errors, warnings = validator.validate_parquet_data(parquet_dir) - - # Report validation statistics - if validator.validation_stats: - console.print("\n[bold]Validation Statistics:[/bold]") - for table_name, stats in validator.validation_stats.items(): - if stats['errors'] == 0: - console.print(f" [green]✓ {table_name}: {stats['valid']}/{stats['sampled']} sampled valid ({stats['total']} total)[/green]") - else: - console.print(f" [red]✗ {table_name}: {stats['valid']}/{stats['sampled']} sampled valid ({stats['errors']} errors)[/red]") - - if errors: - console.print("\n[red]✗ Schema Validation Errors:[/red]") - for error in errors: - console.print(f" [red]• {error}[/red]") - - if warnings: - console.print("\n[yellow]⚠ Schema Validation Warnings:[/yellow]") - for warning in warnings: - console.print(f" [yellow]• {warning}[/yellow]") - - if is_valid and not warnings: - console.print("\n[green]✓ Schema validation passed![/green]") - elif is_valid: - console.print("\n[green]✓ Schema validation passed with warnings[/green]") - - return is_valid - - -if __name__ == "__main__": - import sys - if len(sys.argv) < 2: - print("Usage: python schema_validator.py ") - sys.exit(1) - - product_name = sys.argv[1] - product_dir = Path(f"data/products/{product_name}") - - success = validate_product_schema(product_dir) - sys.exit(0 if success else 1) diff --git a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/run_product_pipeline.sh b/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/run_product_pipeline.sh deleted file mode 100755 index f6c15b72b..000000000 --- a/mpcontribs-lux/mpcontribs/lux/projects/alab/pipelines/run_product_pipeline.sh +++ /dev/null @@ -1,38 +0,0 @@ -#!/bin/bash -# -# Product Pipeline Runner - Uses existing data/venv -# -# Usage: -# ./run_product_pipeline.sh create # Create new product -# ./run_product_pipeline.sh list # List products -# ./run_product_pipeline.sh run --product # Run pipeline (dry run) -# ./run_product_pipeline.sh run --product --upload # Upload (MPContribs + S3) -# ./run_product_pipeline.sh status --product # Check status -# - -set -e - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -cd "$SCRIPT_DIR" - -# Check if venv exists -if [ ! -d "data/venv" ]; then - echo "→ Creating virtual environment..." - python3 -m venv data/venv - echo "✓ Virtual environment created" -fi - -# Activate venv -source data/venv/bin/activate - -# Install/update dependencies -echo "→ Checking dependencies..." -pip install -q -r data/requirements.txt 2>&1 | grep -v "UserWarning" | grep -v "Valid config keys" || true - -# Run the pipeline CLI (default to 'create' if no args) -if [ $# -eq 0 ]; then - python data/pipeline/run_pipeline.py create -else -python data/pipeline/run_pipeline.py "$@" -fi -