WIP NWB auto classes

auto_extensions.py

from pynwb.spec import NWBDatasetSpec, NWBNamespaceBuilder, NWBGroupSpec, NWBAttributeSpec
from pynwb.form.spec import RefSpec
from pynwb import register_class, load_namespaces, NWBFile, NWBHDF5IO
from pynwb.form.utils import docval
from pynwb.file import Subject, NWBContainer, MultiContainerInterface, NWBDataInterface
from pynwb.device import Device



from datetime import datetime
from dateutil.parser import parse as parse_date

import re


name = 'buzsaki_meta'
ns_path = name + ".namespace.yaml"
ext_source = name + ".extensions.yaml"


manipulation = NWBGroupSpec(
    neurodata_type_def='Manipulation',
    neurodata_type_inc='NWBDataInterface',
    quantity='+',
    doc='manipulation',
    attributes=[
        NWBAttributeSpec(name='brain_region_target', dtype='text', doc='Allan Institute Acronym')
    ]
)


virus_injection = NWBGroupSpec(
    neurodata_type_inc='NWBDataInterface',
    neurodata_type_def='VirusInjection', quantity='+',
    doc='notes about surgery that includes virus injection',
    datasets=[NWBDatasetSpec(name='coordinates', doc='(AP, ML, DV) of virus injection',
                             dtype='float', shape=(3,))],
    attributes=[
        NWBAttributeSpec(name='virus', doc='type of virus', dtype='text'),
        NWBAttributeSpec(name='volume', doc='volume of injecting in nL', dtype='float'),
        NWBAttributeSpec(name='rate', doc='rate of injection (nL/s)',
                         dtype='float', required=False),
        NWBAttributeSpec(name='scheme', doc='scheme of injection', dtype='text', required=False),
        NWBAttributeSpec(name='help', doc='help', dtype='text', value='Information about a virus injection')])

virus_injections = NWBGroupSpec(
    neurodata_type_def='VirusInjections',
    neurodata_type_inc='NWBDataInterface',
    name='virus_injections',
    doc='stores virus injections', quantity='?',
    groups=[virus_injection],
    attributes=[
        NWBAttributeSpec(name='help', doc='help', dtype='text', value='Container for virus injections')
    ])

manipulations = NWBGroupSpec(
    neurodata_type_def='Manipulations',
    neurodata_type_inc='NWBDataInterface',
    name='manipulations',
    doc='stores maipulations', quantity='?',
    groups=[manipulation])


surgery = NWBGroupSpec(
    neurodata_type_def='Surgery', doc='information about a specific surgery', quantity='+',
    neurodata_type_inc='NWBDataInterface',
    datasets=[NWBDatasetSpec(name='devices', quantity='?', doc='links to implanted/explanted devices',
                             dtype=RefSpec('Device', 'object'))],
    groups=[virus_injections, manipulations],
    attributes=[
        NWBAttributeSpec(name='start_datetime', doc='datetime in ISO 8601', dtype='text', required=False),
        NWBAttributeSpec(name='end_datetime', doc='datetime in ISO 8601', dtype='text', required=False),
        NWBAttributeSpec(name='weight', required=False, dtype='text',
                         doc='Weight at time of experiment, at time of surgery and at other '
                             'important times'),
        NWBAttributeSpec(name='notes', doc='notes and complications', dtype='text', required=False),
        NWBAttributeSpec(name='anesthesia', doc='anesthesia', dtype='text', required=False),
        NWBAttributeSpec(name='analgesics', doc='analgesics', dtype='text', required=False),
        NWBAttributeSpec(name='antibiotics', doc='antibiotics', dtype='text', required=False),
        NWBAttributeSpec(name='complications', doc='complications', dtype='text', required=False),
        NWBAttributeSpec(name='target_anatomy', doc='target anatomy', dtype='text', required=False),
        NWBAttributeSpec(name='room', doc='place where the surgery took place', dtype='text',
                         required=False),
        NWBAttributeSpec(name='surgery_type', doc='"chronic" or "acute"', dtype='text', required=False),
        NWBAttributeSpec(name='help', doc='help', dtype='text', value='Information about surgery')
    ])

surgeries = NWBGroupSpec(
    neurodata_type_def='Surgeries',
    neurodata_type_inc='NWBDataInterface',
    name='surgeries',
    doc='relevant data for surgeries', quantity='?',
    groups=[surgery],
    attributes=[
        NWBAttributeSpec(name='help', doc='help', dtype='text', value='Container for surgeries')
    ])

histology = NWBGroupSpec(
    neurodata_type_def='Histology',
    neurodata_type_inc='NWBDataInterface',
    name='histology',
    doc='information about histology of subject',
    quantity='?',
    attributes=[
        NWBAttributeSpec(name='file_name', doc='filename of histology images', dtype='text'),
        NWBAttributeSpec(name='file_name_ext', doc='filename extension', dtype='text'),
        NWBAttributeSpec(name='imaging_technique',
                         doc='histology imaging technique (e.g. widefield, confocal, etc.)',
                         dtype='text'),
        NWBAttributeSpec(name='slice_plane', doc='[Coronal, Sagital, Transverse, Other]',
                         required=False, dtype='text'),
        NWBAttributeSpec(name='slice_thickness', doc='thickness of slice (um)', dtype='float',
                         required=False),
        NWBAttributeSpec(name='location_along_axis', doc='Axis orthogal to SlicePlane (mm)',
                         dtype='float', required=False),
        NWBAttributeSpec(name='brain_region_target', doc='Allen Institute acronym',
                         dtype='text', required=False),
        NWBAttributeSpec(name='stainings', doc='stainings', dtype='text', required=False),
        NWBAttributeSpec(name='light_source', doc='wavelength of light source in nm',
                         dtype='float', required=False),
        NWBAttributeSpec(name='image_scale', doc='scale of image (pixels/100um)', dtype='float',
                         required=False),
        NWBAttributeSpec(name='scale_bar', doc='size of image scale bar (um)', dtype='float',
                         required=False),
        NWBAttributeSpec(name='post_processing', doc='[Z-stacked, Stiched]', dtype='text',
                         required=False),
        NWBAttributeSpec(name='user', doc='person involved', dtype='text', required=False),
        NWBAttributeSpec(name='notes', doc='anything else', dtype='text', required=False),
        NWBAttributeSpec(name='help', doc='help', dtype='text', value='Information about Histology')
    ])


subject = NWBGroupSpec(
    neurodata_type_inc='Subject',
    neurodata_type_def='BuzSubject',
    name='subject',
    doc='information about subject',
    groups=[surgeries, histology],
    attributes=[
        NWBAttributeSpec(
            name='sex', required=False, dtype='text',
            doc='Sex of subject. Options: "M": male, "F": female, "O": other, "U": unknown'),
        NWBAttributeSpec(name='species', doc='Species of subject', dtype='text', required=False),
        NWBAttributeSpec(name='strain', dtype='text', doc='strain of animal', required=False),
        NWBAttributeSpec(name='genotype', dtype='text', doc='genetic line of animal', required=False),
        NWBAttributeSpec(name='date_of_birth', dtype='text', doc='in ISO 8601 format', required=False),
        NWBAttributeSpec(name='date_of_death', dtype='text', doc='in ISO 8601 format', required=False),
        NWBAttributeSpec(name='age', doc='age of subject. No specific format enforced.', dtype='text',
                         required=False),
        NWBAttributeSpec(name='gender', dtype='text', required=False,
                         doc='Gender of subject if different from sex.'),
        NWBAttributeSpec(name='earmark', dtype='text', required=False,
                         doc='Earmark of subject'),
        NWBAttributeSpec(name='weight', required=False, dtype='float',
                         doc='Weight at time of experiment, at time of surgery in grams'),
        NWBAttributeSpec(name='help', doc='help', dtype='text', value='Buzsaki subject structure')
    ])

probe = NWBGroupSpec(
    neurodata_type_inc='Device',
    neurodata_type_def='Probe',
    name='probe',
    doc='probe',
    datasets=[
        NWBDatasetSpec(name='coordinates', doc='(AP, ML, DV) of virus injection',
                       dtype='float', shape=(3,)),
        NWBDatasetSpec(name='angles', doc='(degrees) [AP,MD,DV]', dtype='float', shape=(3,))

    ],
    attributes=[
        NWBAttributeSpec(name='nchannels', dtype='int', doc='number of channels'),
        NWBAttributeSpec(name='spike_groups', dtype='int', doc='spike groups'),
        NWBAttributeSpec(name='wire_count', dtype='int', doc='wire count'),
        NWBAttributeSpec(name='write_diameter', dtype='float', doc='diameter of wire'),
        NWBAttributeSpec(name='rotation', dtype='float', doc='rotation of probe'),
        NWBAttributeSpec(name='ground_electrode', dtype='text', doc='e.g. "screw above cerebellum"'),
        NWBAttributeSpec(name='reference_electrode', dtype='text', doc='e.g. "shorted to ground"')
    ]
)

silicon_probe = NWBGroupSpec(
    neurodata_type_inc='Probe',
    neurodata_type_def='SiliconProbe',
    doc='silicon probe',
    attributes=[
        NWBAttributeSpec(name='probe_id', dtype='text', doc='probe id')
    ]
)

tetrode = NWBGroupSpec(
    neurodata_type_inc='Probe',
    neurodata_type_def='Tetrode',
    attributes=[
        NWBAttributeSpec(name='tetrode_count', dtype='int', doc='number of tetrodes')
    ]
)

optical_fiber = NWBGroupSpec(
    neurodata_type_inc='Device',
    neurodata_type_def='OpticalFiber',
    name='OpticalFiber',
    doc='Meta-data about optical fiber',
    attributes=[
        NWBAttributeSpec(name='type', doc='model', dtype='text', required=False),
        NWBAttributeSpec(name='core_diameter', doc='in um', dtype='float', required=False),
        NWBAttributeSpec(name='outer_diameter', doc='in um', dtype='float', required=False),
        NWBAttributeSpec(name='microdrive', doc='whether a microdrive was used (0: not used, 1: used)',
                         dtype='int'),
        NWBAttributeSpec(name='microdrive_lead', doc='um/turn', dtype='float', required=False),
        NWBAttributeSpec(name='microdrive_id', doc='id of microdrive', dtype='int', required=False),
        NWBAttributeSpec(name='help', doc='help', dtype='text', value='Information about optical fiber')
    ]
)

ns_builder = NWBNamespaceBuilder(name + ' extensions', name)

specs = (subject, optical_fiber)
for spec in specs:
    ns_builder.add_spec(ext_source, spec)
ns_builder.export(ns_path)


def obj2docval(spec):

    args_spec = []

    for attrib in spec.attributes:
        if attrib.dtype is 'text':
            _type = str
        else:
            _type = attrib.dtype

        arg_spec = {'name': attrib.name, 'type': _type, 'doc': attrib.doc}
        if not attrib.required:
            arg_spec['default'] = None
        if not attrib.name == 'help':
            args_spec.append(arg_spec)

    for group in spec.groups + spec.datasets:
        arg_spec = {'name': group.name, 'type': group.neurodata_type_def, 'doc': group.doc}
        if group.quantity in ('?', '*'):
            arg_spec['default'] = None

        args_spec.append(arg_spec)

    names = [x['name'] for x in args_spec]
    super_args = eval(spec['neurodata_type_inc']).__init__.__docval__['args']
    for x in super_args:
        if x['name'] not in names:
            args_spec.append(x)

    return tuple(args_spec)


def get_nwbfields(spec):
    vars = [attrib.name for attrib in spec.attributes] + \
           [attrib.name for attrib in spec.datasets] + \
           [{'name': attrib.name, 'child': True} for attrib in spec.groups
            if 'neurodata_type_inc' in attrib or 'neurodata_type_def' in attrib]

    return tuple(vars)

####


load_namespaces(ns_path)

# load custom classes
ns_path = name + '.namespace.yaml'
ext_source = name + '.extensions.yaml'
load_namespaces(ns_path)


def get_class(spec):
    class AutoClass(eval(spec['neurodata_type_inc'])):
        __nwbfields__ = get_nwbfields(spec)

        @docval(*obj2docval(spec))
        def __init__(self, **kwargs):
            super_args = [x['name'] for x in super(AutoClass, self).__init__.__docval__['args']]
            super(AutoClass, self).__init__(**{arg: kwargs[arg] for arg in super_args if arg in kwargs})
            for attr, val in kwargs.items():
                if attr not in super_args:
                    setattr(self, attr, val)
    return AutoClass


def camel2underscore(name):
    s1 = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', name)
    return re.sub('([a-z0-9])([A-Z])', r'\1_\2', s1).lower()


def get_multi_container(spec):
    inner_class_name =spec.groups[0]['neurodata_type_def']
    inner_class = camel2underscore(inner_class_name)
    InnerClass = eval(inner_class_name)

    @register_class(spec['neurodata_type_def'], name)
    class AutoClass(MultiContainerInterface):
        __clsconf__ = {
            'attr': inner_class + 's',
            'type': InnerClass,
            'add': 'add_' + inner_class,
            'get': 'get_' + inner_class,
            'create': 'create_' + inner_class,
        }

        __help = 'container for ' + inner_class + 's'

    return AutoClass


@register_class('Surgery', name)
class Surgery(get_class(surgery)):
    @docval(*obj2docval(surgery))
    def __init__(self, **kwargs):
        super(Surgery, self).__init__(**kwargs)
        if self.surgery_type not in ('chronic', 'acute', None):
            raise ValueError(self.name + ": surgery_type must be 'chronic' or 'acute'")

        if self.start_datetime:
            parse_date(self.start_datetime)

        if self.end_datetime:
            parse_date(self.end_datetime)


@register_class('Surgeries', name)
class Surgeries(get_multi_container(surgeries)):
    pass


@register_class('VirusInjection', name)
class VirusInjection(get_class(virus_injection)):
    pass


@register_class('VirusInjections', name)
class VirusInjections(get_multi_container(virus_injections)):
    pass


@register_class('BuzSubject', name)
class BuzSubject(get_class(subject)):
    @docval(*obj2docval(subject))
    def __init__(self, **kwargs):
        super(BuzSubject, self).__init__(**kwargs)
        if self.sex not in ('M', 'F', 'U'):
            raise ValueError('sex must be M (male), F (female) or U (unknown)')


@register_class('Histology', name)
class Histology(get_class(histology)):
    pass


@register_class('OpticalFiber', name)
class OpticalFiber(get_class(optical_fiber)):
    pass


virus_injections = VirusInjections(source='lab notebook', virus_injections=[
    VirusInjection(
        name='virus_injection1', coordinates=[1., 2., 3.], virus='a', volume=.45,
        source='source', scheme='a')
])

implantation = Surgery(
    name='implantation', notes='test surgery', source='lab notebook',
    virus_injections=virus_injections, anesthesia='a', analgesics='a', antibiotics='a',
    target_anatomy='CA1', room='35C', surgery_type='chronic',
    start_datetime=datetime.utcnow().isoformat() + "Z")
surgeries = Surgeries(source='lab notebook')
surgeries.add_surgery(implantation)

histology = Histology(
    name='histology', source='notebook', file_name='007_histology_files', file_name_ext='png',
    imaging_technique='widefield', slice_plane='Coronal', slice_thickness=100.,
    location_along_axis=21.4, brain_region_target='CA1', stainings='stainings info',
    light_source=300., image_scale=300., scale_bar=100., post_processing='Z-stacked',
    user='name of person', notes='notes')

subject = BuzSubject(
    subject_id='007', genotype='mouse1', species='mouse', age='3 months', weight=20.,
    sex='U', surgeries=surgeries, source='notebook', histology=histology)

nwbfile = NWBFile("source", "a file with metadata", "NB123A", '2018-06-01T00:00:00', subject=subject)

nwbfile.add_device(
    OpticalFiber(
        microdrive=0, source='source', name='optical_fiber1', type='type A',
        core_diameter=.1, outer_diameter=.2, microdrive_lead=2.1, microdrive_id=1))

fname = 'test_ext.nwb'
with NWBHDF5IO(fname, 'w') as io:
    io.write(nwbfile)

with NWBHDF5IO(fname, 'r') as io:
    nwbfile = io.read()
    print(nwbfile.subject.surgeries.surgerys['implantation'].
          virus_injections.virus_injections['virus_injection1'].coordinates[:])

The read step isn't working yet. I chose this approach because it allows you to generate classes automatically but also to alter them, for instance if you wanted to make other methods for the class or add input parsing.

@ajtritt I got it to work!! Let me know what you think!

nice work @bendichter! A couple things:

1. get_nwbfields should create a field with child=True only for group specs with a neurodata_type_inc or neurodata_type_def. This is related to my comment here.
2. In get_class and get_multicontainer, instead of creating of statically creating classes (i.e. class Autoclass(...)), you should dynamically create them using type. This way, you don't have to extend AutoClass just to rename the class.