Subclassing for aind-behavior. How should we organize?

[bruno-f-cruz]

Some generic objects would be nice to subclass in order to more precisely define data structures that we might want to reuse across several different behavioral tasks. An obvious one would be the water calibration.
Under the current data-schema spec, this would be implemented using the Calibration class:

aind-data-schema/src/aind_data_schema/device.py

Line 110 in e92f039

class Calibration(AindModel):

Since this is a generic class aimed at having the flexibility to express several different calibration procedures, it codes both the inputs and the outputs of the calibration as Optional[Dict[str, Any]].

For water calibration, we could come up with a much more precise structure (e.g. inputs would perhaps code delay and water_weight, while outputs could store some sort of model (e.g. just a linear with a slope and offset) or a lookup table. Both could be easily expressed in a schema format.

This raises several questions:

1. Should we PR changes onto the main data-schema repository, or create our "library" where we maintain children of the main repository?
2. Inhering from this parent class might be a bit tricky since the input type is already defined. I wonder if a better approach is to leave the parent Calibration class with generic inputs as Optional[Any]. On the other hand, I believe that any object we set as the input to the children class, should be able to be deserialized back to a Dict[str,Any], since that's what a json object is anyway. So this might not be an issue (?)
3. What would this WaterCalibration object look like? Can you add the JSON you are currently using in the dynamic foraging (@XX-Yin or @hanhou) to store this data?

[XX-Yin]

WaterCalibration.json
@bruno-f-cruz The attached is an example water calibration JSON file.

The meaning of each field is shown in the following code.
WaterCalibrationResults[date_str][valve][valve_open_time][valve_open_interval][cycle].append(total_water)

[bruno-f-cruz]

Ok I am thinking about something along the lines of:

class Calibration(AindModel):
    """Generic calibration class"""

    calibration_date: datetime = Field(..., title="Date and time of calibration")
    device_name: str = Field(..., title="Device name", description="Must match a device name in rig/instrument")
    description: str = Field(..., title="Description", description="Brief description of what is being calibrated")
    input: Optional[Dict[str, Any]] = Field({}, description="Calibration input", title="inputs")
    output: Optional[Dict[str, Any]] = Field({}, description="Calibration output", title="outputs")
    notes: Optional[str] = Field(None, title="Notes")


class WaterCalibrationBolus(AindModel):
    valve_open_interval : float = Field(..., title="Time between two consecutive valve openings (s)", title="Valve open interval", units="s")
    valve_open_time : float = Field(..., title="Valve open interval (s)", title="Valve open time", units="s")
    water_weight: float = Field(..., title="Weight of water delivered (g)", title="Water weight", units="g")
    repeat_count: int = Field(..., title="Number of times the valve opened.", title="Repeat count")


class WaterCalibration(Calibration):
    input: List[WaterCalibrationBolus] = Field(None, title="Input of the calibration")
    output: Union[Dict[str, Any], List[float]] = Field(None, title="Output of the calibration. This can be a list of floats used for a linear model [B0, B1,...] or a generic.")

Wondering how specific we want to be with the model. Is it the case that you are always using a linear fit?

[bruno-f-cruz]

Notes after talking to @hanhou
to the input:

• water_weight should probably be an array where each element corresponds to a measurement of repeat_count trials.

to the output:

• add an average-single-event-weight as a way to code the predicted single bolus volume (i.e. water_weight_array.mean()/repeat_count)
• add quality metrics (r2)
• strongly define slope and offset properties
• add optional property called ValidDomain that defines the range of delays that were used for the calibration calculation

[saskiad]

I actually would prefer we not subclass the Calibration class if possible. I think we want the calibration list to include a calibration of the reward spout, and how you put the data in there should follow a convention and so having something in the ingest script makes sense, but I think adding subclasses into the schema is going to cause bloat.

[bruno-f-cruz]

Would you be open to us starting to compile our own schemas and having units tests to make sure they are compatible with the parents? Or would you rather just maintain everything in the aind-data-schemas repo?

[cindypoo]

we're in need of calibrating multiple instruments, and would be good to find an agreed schema for this. @saskiad do you have guidance on this?

[bruno-f-cruz]

Here's the revised structure. Still needs to go through some testing, but let me know if it makes sense.

from pydantic import PositiveFloat, PositiveInt, Field
from typing import List, Optional, Dict
import aind_data_schema
from aind_data_schema.base import AindModel
from aind_data_schema.models.devices import Calibration
import json

class WaterCalibrationInput(AindModel):
    valve_open_interval: PositiveFloat = Field(..., description="Time between two consecutive valve openings (s)", title="Valve open interval", units="s")
    valve_open_time: PositiveFloat = Field(..., description="Valve open interval (s)", title="Valve open time", units="s")
    water_weight: List[PositiveFloat] = Field(..., description="Weight of water delivered (g)", title="Water weight", units="g", min_length=1)
    repeat_count: int = Field(..., ge=0, description="Number of times the valve opened.", title="Repeat count")

class WaterCalibrationOutput(AindModel):
    interval_average: Optional[Dict[PositiveFloat, PositiveFloat]] = Field(None, description="Dictionary keyed by measured valve interval and corresponding average single event volume.", title="Interval average", units="s")
    slope: float = Field(..., description="Slope of the linear regression : Volume(g) = Slope(g/s) * time(s) + offset(g)", title="Regression slope", units="g/s")
    offset: float = Field(..., description="Offset of the linear regression : Volume(g) = Slope(g/s) * time(s) + offset(g)", title="Regression offset", units="g")
    r2: PositiveFloat = Field(..., description="R2 metric from the linear model.", title="R2", gt=0, le=1)
    valid_domain: Optional[List[PositiveFloat]] = Field(None, description="The optional time-intervals the calibration curve was calculated on.", min_length=2, title="Valid domain", units="s")

class WaterCalibration(Calibration):
    description: Literal["Calibration of the water valve delivery system"] = "Calibration of the water valve delivery system"
    input: List[WaterCalibrationInput] = Field([], title="Input of the calibration") #see note 1.
    output: Optional[WaterCalibrationOutput] = Field(None, title="Output of the calibration.")

Note 1. This assumes the type of the Calibration parent class is changed to AindModel, if not we will need to use a clunkier notation of:
Dict[str: List[WaterCalibrationInput]]

[bruno-f-cruz]

Any thoughts on this @saskiad @dyf ?
I am happy to put a PR together but I want to understand which direction you want to move.

1. Sub classing yes or no?
2. Added classes go to the core aind schema or should we start a separate repo for schemas specific for te behavior team that we compose with core schemas?
3. Are you ok with the proposed water calibration schema?
   Thanks!