Make IO device classes inherit from a common parent to ensure discriminated unions

[bruno-f-cruz]

Is your feature request related to a problem? Please describe.

Consider the following class:

https://github.com/AllenNeuralDynamics/aind-data-schema/blob/728bb06c127978c16a3a4592a26825e1f76c6603/src/aind_data_schema/models/devices.py#L743

It is legitimate to want to pass a Device, DAQChannel as types to this property, since it will allow polymorphic behavior down the line. Unfortunately, there is no way to write a discriminated union since no common properties are shared between the classes. This precludes its proper deserialization as I mentioned before (#518) and it just makes any efforts to ensure such polymorphic deserialization much harder.

Describe the solution you'd like

Make all classes inside Device that we might want to use in cases such as this inherited from the same parent. This parent could have a property that can then be used to build discriminated unions. Could even just be the name of the class tbh

Describe alternatives you've considered Open to suggestions...

Additional context N/A

[jtyoung84]

Describe alternatives you've considered Open to suggestions...

If we made it as a Discriminated Union, then we should at least get errors that the classes are missing the tagged field. Inheriting from a parent class is okay, but as long as they have a shared Literal field, it will also work.

[bruno-f-cruz]

Sure that is another way to go about it, but then we would need to make sure that every single possible class that we might want to include in these unions has that property which I don't think it will be something easy to do without bumping the schema each time. I guess it boils down to how often we think this pattern will occur?

[saskiad]

@bruno-f-cruz can you explain why we need the lick_sensor is a Union to begin with? It came from this issue: https://github.com/AllenNeuralDynamics/aind-data-schema/issues/723 But I still think it's better as just a Device, and then we wouldn't have this issue

[bruno-f-cruz]

I suspect at some point when defining a schema you will need to specify a discriminated union of Union[HarpLickometer, FooLickometer, BarLickometer] as an input type for a field in your mouse platform.

All the aforementioned lickometers should still inherit from device, but you will still want to type hint the parameters of the field to restrict which Devices you can use as lickometers.

To clarify, you don't need the inheritance pattern per se, it is just much easier to maintain the code base if you use it tho. Otherwise you will need to manually add the same field to all lickometers such that you can build discriminated unions later.

[saskiad]

Why do we need separate classes for each lickometer? What would be in those classes?

[saskiad]

in my mind, lick_sensor should just be Device. I'm unclear why it's a union to start with

[bruno-f-cruz]

In the original issue #723 you wrote:

1.

in Device add new Device class (above RewardSpout) called LickDetector with an enum for a field called detector_type with options Capacitive and Piezo Electric. That's the only unique field here

2.

In RewardSpout add an option field called lick_detector that takes the LickDetector class above

Lets take 1) first. If you consider lick_sensor (I will type LickSensor to refer to it as a class) a Device, it will make code difficult to maintain later since you can have different types of lick sensors. HarpLicketlySplit is one, but we have the janelia lickometer too, and I suspect in the building there are other ones available as well. These cluster based on method, but might also use the same method (eg. capacitance) but be completely different devices. So yes, a lick sensor "Device" should be a Device, but different lickometers should likely be different devices too.

2. Now, lick_sensor is also an attribute of RewardSpout. In this case, the attribute needs to be correctly typed. If we assume that all LickSensors are Devices. One could do:

lick_sensor: Device # (could be interesting if you want to let users define their own lickometer devices)
#or
lick_sensor : Union[HarpLickometer, FooLickometer, BarLickometer]

What do you lose by doing this? You can't define a lickometer as being the input pin of a specific device. For instance, imagine that the NIDAQ Channel 0 is connected to some sort of circuit that outputs some sort of IR-beam break near the tongue. If what is recorded is the signal, what do you want to specify as lick_sensor? The IR-beam device (and therefor add it as its own device to the schema) or simply:

lick_sensor: DAQChannel

What is currently in the schemas (see the link in the original post) is a bit of a hack since you are telling users that you allow DaqChannel and Device but you don't afford a discriminated union. A possible solution is to make sure they share a common literal property (for instance via inheritance as proposed) and change the union to a discriminated union.

[saskiad]

why can't these be discriminated by device_type?

[jtyoung84]

In the current implementation:

lick_sensor: Optional[Union[Device, DAQChannel]] 

both Device and DAQChannel are parent classes. The discriminated field needs to be a Literal. We can try adding device_type to the DAQChannel class, and try to make it a Literal field in Device if that's the expected behavior.

[bruno-f-cruz]

Thanks @jtyoung84. Unfortunately, the current problem goes deeper than that. Taking a closer look at the issue, I don't believe that the current pattern is doing what is intended anyway. Since the class doesn't know about its children types that it should deserialize to, the user is losing information on deserialization (see the first part of the example below). Long story short, if a user tries to deserialize from a LicketySplit class, the best it can get is a corresponding Device base class (incidentally, this happens purely by "luck" as, if the two currently assigned classes: DAQChannel and Device were to be compatible, there is no reason why it wouldn't return a DAQChannel either. The next paragraph will expand on this point)

I should also stress that the reason why we even get a Device, to begin with, is due to the default behavior that pydantic provides on deserialization without discriminators.

As I pointed out before, this default behavior should not be relied on, and we should instead be explicit about it using discriminated unions! See #518 and #731 for more details.

Stepping back a bit, we should consider writing down a few code guidelines on how these patterns should be written in the schemas. There should be something, somewhere, that people can reference when they are writing down new classes. Otherwise, these "rogue" patterns will keep popping up in production code.

from pydantic import BaseModel, Field
from typing import Optional, Union, Literal, TypeVar
from aind_data_schema.components import devices as parent
from aind_data_schema_models.harp_types import LicketySplit
from aind_data_schema_models.organizations import Organization
from aind_data_schema.models.organizations import Thorlabs

class Device(parent.Device):
    discriminator: Literal["device"] = "device"

class DAQChannel(parent.DAQChannel):
    discriminator: Literal["daq_channel"] = "daq_channel"

class LickSensor(LicketySplit):
    discriminator: Literal["lick_sensor"] = "lick_sensor"

class Container(BaseModel):
    parent_lick_sensor: Optional[Union[parent.Device, parent.DAQChannel]] = Field(default=None, title="Lick sensor")
    lick_sensor: Optional[Union[Device, DAQChannel, LickSensor]] = Field(
        default=None, title="Lick sensor", discriminator="discriminator"
    )

parent_device = parent.Device(device_type="foo", name="parent_device", serial_number="1234")
parent_daq_channel = parent.DAQChannel(
    channel_name="foo_channel", device_name="parent_device", channel_type=parent.DaqChannelType.AI
)
parent_detector = parent.Detector(
    name="parent_detector",
    manufacturer=Organization.THORLABS,
    data_interface=parent.DataInterface.USB,
    detector_type=parent.DetectorType.CAMERA,
)
parent_lickometer = LicketySplit(device_name="parent_lickometer")

# parent_detector = parent.Detector(
#     name="parent_detector",
#     manufacturer=Thorlabs,
#     data_interface=parent.DataInterface.USB,
#     detector_type=parent.DetectorType.CAMERA,
# ) # This crashes btw. See issue #960 for more details.

TBaseModel = TypeVar("TBaseModel", bound=BaseModel)
def round_trip(container: TBaseModel) -> TBaseModel:
    return container.model_validate_json(container.model_dump_json())

null_container = Container()

for dev in [parent_device, parent_daq_channel, parent_detector]:
    parent_container = Container(parent_lick_sensor=dev)
    deserialized = round_trip(parent_container)
    print(parent_container.parent_lick_sensor == dev)
    print(type(parent_container.parent_lick_sensor))
    print(deserialized.parent_lick_sensor == dev)
    print(type(deserialized.parent_lick_sensor))

# True
# <class 'aind_data_schema.components.devices.Device'>
# True
# <class 'aind_data_schema.components.devices.Device'>
# True
# <class 'aind_data_schema.components.devices.DAQChannel'>
# True
# <class 'aind_data_schema.components.devices.DAQChannel'>
# True
# <class 'aind_data_schema.components.devices.Detector'>
# False
# <class 'aind_data_schema.components.devices.Device'>

### With discriminators
child_device = Device(**parent_device.model_dump())
child_daq_channel = DAQChannel(**parent_daq_channel.model_dump())
child_lick_sensor = LickSensor(**parent_lickometer.model_dump())

for dev in [child_device, child_daq_channel, child_lick_sensor]:
    parent_container = Container(lick_sensor=dev)
    deserialized = round_trip(parent_container)
    print(parent_container.lick_sensor == dev)
    print(type(parent_container.lick_sensor))
    print(deserialized.lick_sensor == dev)
    print(type(deserialized.lick_sensor))

# True
# <class '__main__.Device'>
# True
# <class '__main__.Device'>
# True
# <class '__main__.DAQChannel'>
# True
# <class '__main__.DAQChannel'>
# True
# <class '__main__.LickSensor'>
# True
# <class '__main__.LickSensor'>

[saskiad]

@bruno-f-cruz I think we are using the lick_sensor in two different ways and I wonder if instead we want to disambiguate them. The lick_sensor was originally to indicate what kind of sensor is used, such as the piezoelectric sensor. This is connected to a daq, and that daq should be listed among the daq devices with the sensor listed as the device attached to it. But I don't know that we want to list the daq device here. In that case, I'd actually propose we remove the daq_device from this and only have it as a device.

[bruno-f-cruz]

Do you mind writing down the type signature you are thinking about? I can't tell what you mean by:

Remove the daq_device from this and only have it as a device

There is no daq_device type/field here. Do you mean DaqChannel? If so, would people be ok with simply listing down what devices can accept lickometers and not specifying which pin is it connected to? If that's ok, then sure, have an upper bound on device and discrimination should be easy, I agree!

However, I think this pattern is worth addressing more generally, as it was the original intent of this issue.

[saskiad]

Yes, I meant DaqChannel. I propose this:

lick_sensor: Optional[Device] = Field(default=None, title="Lick sensor")

I actually think we want to avoid situations where things can be DAQ or Devices and keep the concepts separate... but happy to look at other instances where it might make more sense.

[bruno-f-cruz]

As I pointed out in https://github.com/AllenNeuralDynamics/aind-data-schema/issues/788#issuecomment-2156662579, that signature will lead to lossy deserialization of the Device, currently. If you are ok with that, and this information is to be captured elsewhere, it may be ok.