Creating a new model

This tutorial describes the steps necessary to define a new model type using jwst.datamodels.

For further reading and details, see the reference materials in Metadata.

In this tutorial, we’ll go through the process of creating a new type of model for a file format used for storing the bad pixel mask for JWST’s MIRI instrument. This file format has a 2D array containing a bit field for each of the pixels, and a table describing what each of the bits in the array means.

Note

While an attempt is made to present a real-world example here, it may not reflect the actual final format of this file type, which is still subject to change at the time of this writing.

This example will be built as a third-party Python package, i.e. not part of jwst.datamodels itself. Doing so adds a few extra wrinkles to the process, and it’s most helpful to show what those wrinkles are. To skip ahead and just see the example in its entirety, see the examples/custom_model directory within the jwst.datamodels source tree.

Directory layout

The bare minimum directory layout for a Python package that creates a custom model is as below:

.
|-- lib
|   |--- __init__.py
|   |--- bad_pixel_mask.py
|   |--- schemas
|   |--- bad_pixel_mask.schema.yaml
|   |--- tests
|       |--- __init__.py
|       |--- test_bad_pixel_mask.py
|       |--- data
|       |--- bad_pixel_mask.fits
|--- setup.py

The main pieces are the new schema in bad_pixel_mask.schema.yaml, the custom model class in bad_pixel_mask.py, a setup.py file to install the package, and some unit tests and associated data. Normally, you would also have some code that uses the custom model included in the package, but that isn’t included in this minimal example.

The schema file

Let’s start with the schema file, bad_pixel_mask.schema.yaml. There are a few things it needs to do:

  1. It should contain all of the core metadata from the core schema that ships with jwst.datamodels. In JSON Schema parlance, this schema “extends” the core schema. In object-oriented programming terminology, this could be said that our schema “inherits from” the core schema. It’s all the same thing.

  2. Define the pixel array containing the information about each of the bad pixels. This will be an integer for each pixel where each bit is ascribed a particular meaning.

  3. Define a table describing what each of the bit fields in the pixel array means. This will have three columns: one for the bit field’s number (a power of 2), one for a name token to identify it, and one with a human-readable description.

At the top level, every JSON schema must be a mapping (dictionary) of type “object”, and should include the core schema:

allOf:
   - $ref: "http://jwst.stsci.edu/schemas/core.schema.yaml"
   - type: object
     properties:
        ...

There’s a lot going on in this one item. $ref declares the schema fragment that we want to include (the “base class” schema). Here, the $ref mapping causes the system to go out and fetch the content at the given URL, and then replace the mapping with that content.

The $ref URL can be a relative URL, in which case it is relative to the schema file where $ref is used. In our case, however, it’s an absolute URL. Before you visit that URL to see what’s there, I’ll save you the trouble: there is nothing at that HTTP address. The host jwst.stsci.edu is recognized as a “special” address by the system that causes the schema to be looked up alongside installed Python code. For example, to refer to a (hypothetical) my_instrument schema that ships with a Python package called astroboy, use the following URL:

http://jwst.stsci.edu/schemas/astroboy/my_instrument.schema.yaml

The “package” portion may be omitted to refer to schemas in the jwst.datamodels core, which is how we arrive at the URL we’re using here:

http://jwst.stsci.edu/schemas/core.schema.yaml

Note

At some time in the future, we will actually be hosting schemas at a URL similar to the one above. This will allow schemas to be shared with tools built in languages other than Python. Until we have that hosting established, this works quite well and does not require any coordination among Python packages that define new models. Keep an eye out if you use this feature, though – the precise URL used may change.

The next part of the file describes the array data, that is, things that are Numpy arrays on the Python side and images or tables on the FITS side.

First, we describe the main "dq" array. It’s declared to be 2-dimensional, and each element is an unsigned 32-bit integer:

properties:
  dq:
    title: Bad pixel mask
    fits_hdu: DQ
    default: 0
    ndim: 2
    datatype: uint16

The next entry describes a table that will store the mapping between bit fields and their meanings. This table has four columns:

  • BIT: The value of the bit field (a power of 2)

  • VALUE: The value resulting when raising 2 to the BIT power

  • NAME: The name used to refer to the bit field

  • DESCRIPTION: A longer, human-readable description of the bit field

dq_def:
  title: DQ flag definitions
  fits_hdu: DQ_DEF
  dtype:
    - name: BIT
      datatype: uint32
    - name: VALUE
      datatype: uint32
    - name: NAME
      datatype: [ascii, 40]
    - name: DESCRIPTION
      datatype: [ascii, 80]

And finally, we add a metadata element that is specific to this format. To avoid recomputing it repeatedly, we’d like to store a sum of all of the “bad” (i.e. non-zero) pixels stored in the bad pixel mask array. In the model, we want to refer to this value as model.meta.bad_pixel_count. In the FITS file, lets store this in the primary header in a keyword named BPCOUNT:

meta:
  properties:
    bad_pixel_count:
      type: integer
      title: Total count of all bad pixels
      fits_keyword: BPCOUNT

That’s all there is to the schema file, and that’s the hardest part.

The model class

Now, let’s see how this schema is tied in with a new Python class for the model.

First, we need to import the JwstDataModel class, which is the base class for all models:

from stdatamodels.jwst.datamodels import JwstDataModel

Then we create a new Python class that inherits from JwstDataModel, and set its schema_url class member to point to the schema that we just defined above:

class MiriBadPixelMaskModel(JwstDataModel):
    schema_url = "bad_pixel_mask.schema.yaml"

Here, the schema_url has all of the “magical” URL abilities described above when we used the $ref feature. However, here we are using a relative URL. In this case, it is relative to the file in which this class is defined, with a small twist to avoid intermingling Python code and schema files: It looks for the given file in a directory called schemas inside the directory containing the Python module in which the class is defined.

As an alternative, we could just as easily have said that we want to use the image schema from the core without defining any extra elements, by setting schema_url to:

schema_url = "http://jwst.stsci.edu/schemas/image.schema.yaml"

Note

At this point you may be wondering why both the schema and the class have to inherit from base classes. Certainly, it would have been more convenient to have the inheritance on the Python side automatically create the inheritance on the schema side (or vice versa). The reason we can’t is that the schema files are designed to be language-agnostic: it is possible to use them from an entirely different implementation of the jwst.datamodels framework possibly even written in a language other than Python. So the schemas need to “stand alone” from the Python classes. It’s certainly possible to have the schema inherit from one thing and the Python class inherit from another, and the jwst.datamodels framework won’t and can’t really complain, but doing that is only going to lead to confusion, so just don’t do it.

Within this class, we’ll define a constructor. All model constructors must take the highly polymorphic init value as the first argument. This can be a file, another model, or all kinds of other things. See the docstring of jwst.datamodels.JwstDataModel.__init__ for more information. But we’re going to let the base class handle that anyway.

The rest of the arguments are up to you, but generally it’s handy to add a couple of keyword arguments so the user can data arrays when creating a model from scratch. If you don’t need to do that, then technically writing a new constructor for the model is optional:

def __init__(self, init=None, dq=None, dq_def=None, **kwargs):
    """
    A data model to represent MIRI bad pixel masks.

    Parameters
    ----------
    init : any
        Any of the initializers supported by `~jwst.datamodels.JwstDataModel`.

    dq : numpy array
        The data quality array.

    dq_def : numpy array
        The data quality definitions table.
    """
    super(MiriBadPixelMaskModel, self).__init__(init=init, **kwargs)

    if dq is not None:
        self.dq = dq

    if dq_def is not None:
        self.dq_def = dq_def

The super.. line is just the standard Python way of calling the constructor of the base class. The rest of the constructor sets the arrays on the object if any were provided.

The other methods of your class may provide additional conveniences on top of the underlying file format. This is completely optional and if your file format is supported well enough by the underlying schema alone, it may not be necessary to define any extra methods.

In the case of our example, it would be nice to have a function that, given the name of a bit field, would return a new array that is True wherever that bit field is true in the main mask array. Since the order and content of the bit fields are defined in the dq_def table, the function should use it in order to do this work:

def get_mask_for_field(self, name):
    """
    Returns an array that is `True` everywhere a given bitfield is
    True in the mask.

    Parameters
    ----------
    name : str
        The name of the bit field to retrieve

    Returns
    -------
    array : boolean numpy array
        `True` everywhere the requested bitfield is `True`.  This
        is the same shape as the mask array.  This array is a copy
        and changes to it will not affect the underlying model.
    """
    # Find the field value that corresponds to the given name
    field_value = None
    for value, field_name, title in self.dq_def:
        if field_name == name:
            field_value = value
            break
    if field_value is None:
        raise ValueError("Field name {0} not found".format(name))

    # Create an array that is `True` only for the requested
    # bit field
    return self.dq & field_value

One thing to note here: this array is semantically a “copy” of the underlying data. Most Numpy arrays in the model framework are mutable, and we expect that changing their values will update the model itself, and be saved out by subsequent saves to disk. Since the array we are returning here has no connection back to the model’s main data array (mask), it’s helpful to remind the user of that in the docstring, and not present it as a member or property, but as a getter function.

Note

Since handling bit fields like this is such a commonly useful thing, it’s possible that this functionality will become a part of jwst.datamodels itself in the future. However, this still stands as a good example of something someone may want to do in a custom model class.

Lastly, remember the meta.bad_pixel_count element we defined above? We need some way to make sure that whenever the file is written out that it has the correct value. The model may have been loaded and modified. For this, JwstDataModel has the on_save method hook, which may be overridden by the subclass to add anything that should happen just before saving:

def on_save(self, path):
    super(MiriBadPixelMaskModel, self).on_save(path)

    self.meta.bad_pixel_count = np.sum(self.mask != 0)

Note that here, like in the constructor, it is important to “chain up” to the base class so that any things that the base class wants to do right before saving also happen.

The setup.py script

Writing a setup.py script is beyond the scope of this tutorial but it’s worth noting one thing. Since the schema files are not Python files, they are not automatically picked up by setuptools, and must be included in the package_data option. A complete, yet minimal, setup.py is presented below:

#!/usr/bin/env python

from setuptools import setup

setup(
    name='custom_model',
    description='Custom model example for jwst.datamodels',
    packages=['custom_model', 'custom_model.tests'],
    package_dir={'custom_model': 'lib'},
    package_data={'custom_model': ['schemas/*.schema.yaml'],
                  'custom_model.tests' : ['data/*.fits']}
    )

Using the new model

The new model can now be used. For example, to get the locations of all of the “hot” pixels:

from custom_model.bad_pixel_mask import MiriBadPixelMaskModel

with MiriBadPixelMaskModel("bad_pixel_mask.fits") as dm:
    hot_pixels = dm.get_mask_for_field('HOT')

A table-based model

In addition to n-dimensional data arrays, models can also contain tabular data. For example, the photometric correction reference file used in the JWST calibration pipeline consists of a table with several columns. The schema file for one of these models looks like this:

title: NIRISS SOSS photometric flux conversion data model
allOf:
  - $ref: "referencefile.schema.yaml"
  - $ref: "keyword_exptype.schema.yaml"
  - $ref: "keyword_pexptype.schema.yaml"
  - $ref: "keyword_pixelarea.schema.yaml"
  - type: object
    properties:
      phot_table:
        title: Photometric flux conversion factors table
        fits_hdu: PHOTOM
        datatype:
          - name: filter
            datatype: [ascii, 12]
          - name: pupil
            datatype: [ascii, 15]
          - name: order
            datatype: int16
          - name: photmj
            datatype: float32
          - name: uncertainty
            datatype: float32
          - name: nelem
            datatype: int16
          - name: wavelength
            datatype: float32
            ndim: 1
          - name: relresponse
            datatype: float32
            ndim: 1
          - name: reluncertainty
            datatype: float32
            ndim: 1

In this particular table the first 6 columns contain scalar entries of types string, float, and integer. The entries in the final 3 columns, on the other hand, contain 1-D float arrays (vectors). The “ndim” attribute is used to specify the number of dimensions the arrays are allowed to have.

The corresponding python module containing the data model class is quite simple:

class NisSossPhotomModel(ReferenceFileModel):
    """
    A data model for NIRISS SOSS photom reference files.
    """
    schema_url = "nissoss_photom.schema"