Source code for allensdk.brain_observatory.behavior.write_nwb.ophys.schemas

import marshmallow as mm
import pandas as pd
from allensdk.brain_observatory.behavior.behavior_project_cache.project_apis.data_io.behavior_project_cloud_api import (  # noqa: E501
    sanitize_data_columns,
)
from allensdk.brain_observatory.behavior.behavior_project_cache.tables.metadata_table_schemas import (  # noqa: E501
    OphysExperimentMetadataSchema,
)
from allensdk.brain_observatory.behavior.write_nwb.behavior.schemas import (
    BaseInputSchema,
    RaisingSchema,
)
from argschema.fields import Int, List, Nested, OutputFile




class OphysExperimentInputSchema(BaseInputSchema):
    ophys_experiment_id = Int(
        required=True, description="Id of OphysExperiment to create."
    )

    ophys_container_experiment_ids = List(
        Int,
        required=False,
        cli_as_single_argument=True,
        description="Subset of the experiment ids in the same container to be "
        "released. Experiment Ids are pulled from input metadata "
        "table. Useful for when experiments are excluded from the "
        "release and certain summary values (e.g. "
        "targeted_imaging_depth) must be recalculated from the "
        "released experiments.",
        default=[],
    )

    ophys_experiment_metadata = Nested(
        OphysExperimentMetadataSchema,
        required=True,
        description="Data pertaining to an ophys experiment.",
    )



    @mm.pre_load
    def retreive_metadata(self, data, **kwargs):
        """Load the data from csv using Pandas the same as the
        project_cloud api.
        """
        df = sanitize_data_columns(
            data["metadata_table"], dtype_convert={"mouse_id": str}
        )
        df.set_index("ophys_experiment_id", inplace=True)
        try:
            # Enforce type as we haven't enfoced type in the
            # schema yet.
            oe_row = df.loc[int(data["ophys_experiment_id"])]
        except KeyError:
            raise KeyError(
                f"Ophys experiment id {data['ophys_experiment_id']} "
                "not in input OphysExperimentTable. Exiting."
            )

        data["ophys_experiment_metadata"] = self._get_behavior_metadata(oe_row)
        data["ophys_experiment_metadata"]["behavior_session_id"] = oe_row[
            "behavior_session_id"
        ]

        # Ophys Experiment specific data.
        data["ophys_experiment_metadata"]["imaging_depth"] = oe_row[
            "imaging_depth"
        ]
        imaging_plane_group = oe_row["imaging_plane_group"]
        if pd.isna(imaging_plane_group):
            imaging_plane_group = None
        data["ophys_experiment_metadata"][
            "imaging_plane_group"
        ] = imaging_plane_group
        data["ophys_experiment_metadata"]["indicator"] = oe_row["indicator"]
        data["ophys_experiment_metadata"]["ophys_container_id"] = oe_row[
            "ophys_container_id"
        ]
        data["ophys_experiment_metadata"]["ophys_experiment_id"] = oe_row.name
        data["ophys_experiment_metadata"]["ophys_session_id"] = oe_row[
            "ophys_session_id"
        ]
        data["ophys_experiment_metadata"]["targeted_imaging_depth"] = oe_row[
            "targeted_imaging_depth"
        ]
        data["ophys_experiment_metadata"]["targeted_structure"] = oe_row[
            "targeted_structure"
        ]

        data["ophys_container_experiment_ids"] = df[
            df["ophys_container_id"] == oe_row["ophys_container_id"]
        ].index.tolist()

        return data






class OutputSchema(RaisingSchema):
    input_parameters = Nested(OphysExperimentInputSchema)
    output_path = OutputFile(
        required=True,
        description="Path of output NWB file.",
    )