How to use the asdf.tags.transform.basic.TransformType function in asdf

def assert_equal(cls, a, b):
        from astropy.modeling import mappings
        # TODO: If models become comparable themselves, remove this.
        TransformType.assert_equal(a, b)
        assert (isinstance(a, mappings.Identity) and
                isinstance(b, mappings.Identity) and
                a.n_inputs == b.n_inputs)

def to_tree_transform(cls, model, ctx):
        node = {'factor': model.factor.value}
        return yamlutil.custom_tree_to_tagged_tree(node, ctx)

    @classmethod
    def assert_equal(cls, a, b):
        from astropy import modeling

        # TODO: If models become comparable themselves, remove this.
        TransformType.assert_equal(a, b)
        assert (isinstance(a, modeling.models.Scale) and
                isinstance(b, modeling.models.Scale))
        assert_array_equal(a.factor, b.factor)


class PolynomialType(TransformType):
    name = "transform/polynomial"
    types = ['astropy.modeling.models.Polynomial1D',
             'astropy.modeling.models.Polynomial2D']

    @classmethod
    def from_tree_transform(cls, node, ctx):
        from astropy import modeling

        coefficients = np.asarray(node['coefficients'])
        n_dim = coefficients.ndim

        if n_dim == 1:
            model = modeling.models.Polynomial1D(coefficients.size - 1)
            model.parameters = coefficients
        elif n_dim == 2:
            shape = coefficients.shape

def assert_equal(cls, a, b):
        from astropy import modeling

        # TODO: If models become comparable themselves, remove this.
        TransformType.assert_equal(a, b)
        assert (isinstance(a, modeling.models.Shift) and
                isinstance(b, modeling.models.Shift))
        assert_array_equal(a.offset.value, b.offset.value)

    @classmethod
    def to_tree_transform(cls, model, ctx):
        node = {'matrix': model.matrix.value, 'translation': model.translation.value}
        return yamlutil.custom_tree_to_tagged_tree(node, ctx)

    @classmethod
    def assert_equal(cls, a, b):
        # TODO: If models become comparable themselves, remove this.
        TransformType.assert_equal(a, b)
        assert (a.__class__ == b.__class__)
        assert_array_equal(a.matrix, b.matrix)
        assert_array_equal(a.translation, b.translation)


class Rotate2DType(TransformType):
    name = "transform/rotate2d"
    types = ['astropy.modeling.rotations.Rotation2D']

    @classmethod
    def from_tree_transform(cls, node, ctx):
        from astropy import modeling

        return modeling.rotations.Rotation2D(node['angle'])

    @classmethod
    def to_tree_transform(cls, model, ctx):
        return {'angle': model.angle.value}

    @classmethod
    def assert_equal(cls, a, b):
        from astropy import modeling

def assert_equal(cls, a, b):
        # TODO: If models become comparable themselves, remove this.
        TransformType.assert_equal(a, b)
        from ...tests.helpers import assert_tree_match
        assert_tree_match(a._tree.left.value, b._tree.left.value)
        assert_tree_match(a._tree.right.value, b._tree.right.value)
        assert a._tree.value == b._tree.value

'&'  : 'concatenate'
}


_tag_to_method_mapping = {
    'add'         : '__add__',
    'subtract'    : '__sub__',
    'multiply'    : '__mul__',
    'divide'      : '__truediv__',
    'power'       : '__pow__',
    'compose'     : '__or__',
    'concatenate' : '__and__'
}


class CompoundType(TransformType):
    name = ['transform/' + x for x in _tag_to_method_mapping.keys()]
    types = ['astropy.modeling.core._CompoundModel']
    handle_dynamic_subclasses = True

    @classmethod
    def from_tree_tagged(cls, node, ctx):
        from astropy import modeling

        tag = node._tag[node._tag.rfind('/')+1:]
        tag = tag[:tag.rfind('-')]

        oper = _tag_to_method_mapping[tag]
        left = yamlutil.tagged_tree_to_custom_tree(
            node['forward'][0], ctx)
        if not isinstance(left, modeling.Model):
            raise TypeError("Unknown model type '{0}'".format(

    @classmethod
    def to_tree(cls, data, ctx):
        return data


# TODO: This is just here for bootstrapping and will go away eventually
if HAS_ASTROPY:
    class GenericModel(mappings.Mapping):
        def __init__(self, n_inputs, n_outputs):
            mapping = tuple(range(n_inputs))
            super(GenericModel, self).__init__(mapping)
            self._outputs = tuple('x' + str(idx) for idx in range(self.n_outputs + 1))


class GenericType(TransformType):
    name = "transform/generic"
    if HAS_ASTROPY:
        types = [GenericModel]

    @classmethod
    def from_tree_transform(cls, node, ctx):
        return GenericModel(
            node['n_inputs'], node['n_outputs'])

    @classmethod
    def to_tree_transform(cls, data, ctx):
        return {
            'n_inputs': data.n_inputs,
            'n_outputs': data.n_outputs
        }

# -*- coding: utf-8 -*-

from __future__ import absolute_import, division, unicode_literals, print_function

import numpy as np
from numpy.testing import assert_array_equal

from ... import yamlutil

from .basic import TransformType


__all__ = ['ShiftType', 'ScaleType', 'PolynomialType']


class ShiftType(TransformType):
    name = "transform/shift"
    types = ['astropy.modeling.models.Shift']

    @classmethod
    def from_tree_transform(cls, node, ctx):
        from astropy import modeling

        offset = node['offset']
        if not np.isscalar(offset):
            raise NotImplementedError(
                "Asdf currently only supports scalar inputs to Shift transform.")

        return modeling.models.Shift(offset)

    @classmethod
    def to_tree_transform(cls, model, ctx):

try:
    import astropy
except ImportError:
    HAS_ASTROPY = False
else:
    HAS_ASTROPY = True
    from astropy.utils import minversion
    ASTROPY_13 = minversion(astropy, "1.3.dev16506")

if HAS_ASTROPY and ASTROPY_13:
    __all__ = ['TabularType']
else:
    __all__ = []


class TabularType(TransformType):
    name = "transform/tabular"
    if HAS_ASTROPY and ASTROPY_13:
        types = [astropy.modeling.models.Tabular2D,
                 astropy.modeling.models.Tabular1D
                ]
    else:
        types = []

    @classmethod
    def from_tree_transform(cls, node, ctx):
        from astropy import modeling
        lookup_table = node.pop("lookup_table")
        dim = lookup_table.ndim
        name = node.get('name', None)
        fill_value = node.pop("fill_value", None)
        if dim == 1:

def assert_equal(cls, a, b):
        # TODO: If models become comparable themselves, remove this.
        TransformType.assert_equal(a, b)
        assert a.__class__ == b.__class__