User-defined features

[kecnry]

This PR adds support for defining custom features in-line that can either modify models after they are computed (but before they are fitted) or modify component mesh quantities.

Current restrictions:

• Inheritance from other features is ok, but super is not allowed and will result in an error.
• Anything needed within a defined method must be manually imported (or possibly assigned in the init).
• Overwriting the code after attaching is possible, but changing parameterization requires removing and adding a new feature into the bundle.

DatasetFeature:

class CustomFeature(DatasetFeature):
    allowed_dataset_kinds = ['lc', 'rv', 'lp']

    @classmethod
    def create_feature_parameters(self, feature, **kwargs):
        ...
        return PS(params), list_of_constraints

    @classmethod
    def parse_bundle(cls, b, feature_ps):
        return {}

    @classmethod
    def run_checks_compute(cls, b, feature_ps, compute_ps):
        return [{}]

    def modify_data_for_estimators(self, b, data_ps, **data_arrays):
        """
        Modify the data parameters for the estimators.
        This is called before the data is passed to the estimators.
        """
        return {}

    def modify_model(self, b, model_ps):
        return None

ComponentFeature:

class CustomFeature(ComponentFeature):
    """
    Note that for all features, each of the methods below will be called.  So
    changing the coordinates WILL affect the original/intrinsic loggs which
    will then be used as input for that method call.

    In other words, its probably safest if each feature only overrides a
    SINGLE one of the methods.  Overriding multiple methods should be done
    with great care.

    Each feature may or may not require recomputing a mesh, depending on the
    kind of change it exacts to the mesh. For example, pulsations will require
    recomputing a mesh while spots will not. By default, the mesh will be
    recomputed (set in this superclass' `__init__()` method) but inherited
    classes should overload `self.remeshing_required`.
    """
    allowed_component_kinds = ['star', 'envelope', 'orbit']

    @classmethod
    def create_feature_parameters(self, feature, **kwargs):
        ...
        return PS(params), list_of_constraints

    @classmethod
    def parse_bundle(cls, b, feature_ps):
        return {}

    @classmethod
    def run_checks_compute(cls, b, feature_ps, compute_ps):
        return [{}]

    def requires_remeshing(self):
        return True

    def modify_coords_for_computations(self, coords_for_computations, s, t):
        """
        Method for a feature to modify the coordinates.  Coordinates are
        modified AFTER scaling but BEFORE being placed in orbit.

        NOTE: coords_for_computations affect physical properties only and
        not geometric properties (areas, eclipse detection, etc).  If you
        want to override geometric properties, use the hook for
        modify_coords_for_observations as well.

        Features that affect coordinates_for_computations should override
        this method
        """
        return coords_for_computations

    def modify_coords_for_observations(self, coords_for_computations, coords_for_observations, s, t):
        """
        Method for a feature to modify the coordinates.  Coordinates are
        modified AFTER scaling but BEFORE being placed in orbit.

        NOTE: coords_for_observations affect the geometry only (areas of each
        element and eclipse detection) but WILL NOT affect any physical
        parameters (loggs, teffs, intensities).  If you want to override
        physical parameters, use the hook for modify_coords_for_computations
        as well.

        Features that affect coordinates_for_observations should override this method.
        """
        return coords_for_observations

    def cartesian_to_spherical(self, roche_coords):
        return r, colat, longitude

    def modify_rvs(self, rvs, orbit_vel, roche_coords, s=[0., 0., 1.], t=None):
        """
        Method for a feature to modify the radial velocities.

        Features that affect radial velocities (RV+LP datasets) should override this method

        NOTE: orbit_vel[2] is in the OPPOSITE direction of the radial velocity
        """
        return rvs

    def modify_loggs(self, loggs, roche_coords, s=[0., 0., 1.], t=None):
        """
        Method for a feature to modify the loggs.

        Features that affect loggs should override this method
        """
        return loggs

    def modify_teffs(self, teffs, roche_coords, s=[0., 0., 1.], t=None):
        """
        Method for a feature to modify the teffs.

        Features that affect teffs should override this method
        """
        return teffs

    def modify_intensities(self, abs_normal_intensities, abs_intensities,
                           mus, pblum_scale, extinct_factors, boost_factors,
                           roche_coords, s=[0., 0., 1.], t=None):
        """
        Method for a feature to modify the intensities.
        Features that affect intensities should override this method

        Arguments
        ----------
        * `abs_normal_intensities` (ndarray): Absolute normal intensities, already multiplied
            by `extinct_factors`.
        * `abs_intensities` (ndarray): Absolute projected intensities, already multiplied by
            `extinct_factors` and `boost_factors`.
        * `mus` (ndarray): Cosine of the angle between the normal vector and the line of sight
        * `pblum_scale` (ndarray): Scale factor for the pblum, that will be applied to the abs_intensities
            AFTER modify_intensities to result in the scaled intensities
        * `extinct_factors` (ndarray): Extinction factors for the intensities, already applied
        * `boost_factors` (ndarray): Boost factors for the intensities, already applied
        * `roche_coords` (ndarray): Roche coordinates for the computations
        * `s` (array-like): Spin vector in Roche coordinates
        * `t` (float): Current time
        """
        return abs_normal_intensities, abs_intensities

Example of a DatasetFeature (SinusoidalThirdLight):

import phoebe
from phoebe.features import DatasetFeature
from phoebe.parameters import FloatParameter, ParameterSet, constraint
from astropy import units as u
import numpy as np

b = phoebe.default_binary()
b.add_dataset('lc', compute_phases=phoebe.linspace(0,1,101))
b.set_value('irrad_method', 'none')
b.set_value_all('distortion_method', 'sphere')
b.run_compute(model='no_features')

class SinusoidalThirdLight(DatasetFeature):
    allowed_dataset_kinds = ['lc']

    @classmethod
    def create_feature_parameters(self, feature, **kwargs):
        params = []
        params += [FloatParameter(qualifier='amplitude',
                                  latexfmt=r'A_\mathrm{{ {feature} }}',
                                  value=kwargs.get('amplitude', 1.0),
                                  default_unit=u.W/u.m**2,
                                  description='Amplitude of the third light sinusoidal contribution')]
        params += [FloatParameter(qualifier='period',
                                  latexfmt=r'P_\mathrm{{ {feature} }}',
                                  value=kwargs.get('period', 1.0),
                                  default_unit=u.d,
                                  description='Period of the third light sinusoidal contribution')]
        params += [FloatParameter(qualifier='freq',
                                  latexfmt=r'f_\mathrm{{ {feature} }}',
                                  value=kwargs.get('freq', 2*np.pi/3.0),
                                  default_unit=u.rad/u.d, advanced=True,
                                  description='Frequency of the third light sinusoidal contribution')]

        constraints = [(constraint.freq, feature, 'feature')]

        return ParameterSet(params), constraints

    @classmethod
    def parse_bundle(cls, b, feature_ps):
        t0 = b.get_value(qualifier='t0', context='system', unit='d')
        feature_dict = cls.parse_from_feature_ps(b, feature_ps,
                                                 [{'qualifier': 'period', 'unit': 'd'},
                                                  {'qualifier': 'amplitude', 'unit': 'W/m**2'}])
        return dict(t0=t0, **feature_dict)
    
    def modify_model(self, b, model_ps):
        from astropy import units as u
        import numpy as np
        _skip_filter_checks = {'check_default': False, 'check_visible': False}

        for flux_param in model_ps.filter(qualifier='fluxes', **_skip_filter_checks).tolist():
            times = model_ps.get_value(qualifier='times', dataset=flux_param.dataset, unit=u.d, **_skip_filter_checks)
            flux_param.set_value(flux_param.get_value() + self.kwargs['amplitude'] * np.sin(2 * np.pi * (times - self.kwargs['t0']) / self.kwargs['period']), ignore_readonly=True, **_skip_filter_checks)

b.add_feature(SinusoidalThirdLight, dataset='lc01', amplitude=0.5, feature='sin_l3')
#b.add_dataset('mesh', compute_phases=phoebe.linspace(0,1,51), columns='intensities*')
b.run_compute(model='with_sin_l3')

#b.plot(kind='mesh', animate=True, fc='intensities@lc01', ec='face', save='./mesh_with_sin_l3.gif')
b.plot(kind='lc', show=True, legend=True)

Example of a ComponentFeature (SinusoidalIntensities):

import phoebe
from phoebe.features import ComponentFeature
from phoebe.parameters import FloatParameter, ParameterSet, constraint
from astropy import units as u
import numpy as np

b = phoebe.default_binary()
b.add_dataset('lc', compute_phases=phoebe.linspace(0,1,101))
b.set_value('irrad_method', 'none')
b.set_value_all('distortion_method', 'sphere')

b.run_compute(model='no_features')

class SinusoidalIntensities(ComponentFeature):
    @classmethod
    def create_feature_parameters(self, feature, **kwargs):
        params = []
        params += [FloatParameter(qualifier='amplitude',
                                  latexfmt=r'A_\mathrm{{ {feature} }}',
                                  value=kwargs.get('amplitude', 0.1),
                                  default_unit=u.dimensionless_unscaled,
                                  description='Relative amplitude of the sinusoidal intensities contribution')]
        params += [FloatParameter(qualifier='period',
                                  latexfmt=r'P_\mathrm{{ {feature} }}',
                                  value=kwargs.get('period', 1.0),
                                  default_unit=u.d,
                                  description='Period of the sinusoidal intensities contribution')]
        params += [FloatParameter(qualifier='freq',
                                  latexfmt=r'f_\mathrm{{ {feature} }}',
                                  value=kwargs.get('freq', 2*np.pi/3.0),
                                  default_unit=u.rad/u.d, advanced=True,
                                  description='Frequency of the sinusoidal intensities contribution')]

        constraints = [(constraint.freq, feature, 'feature')]

        return ParameterSet(params), constraints
    
    @classmethod
    def parse_bundle(cls, b, feature):
        t0 = b.get_value(qualifier='t0', context='system', unit=u.d, **_skip_filter_checks)
        feature_dict =  cls.parse_from_feature_ps(b, feature,
                                                  ['amplitude',
                                                  {'qualifier': 'period', 'unit': 'd'}])
        return dict(t0=t0, **feature_dict)

    def modify_intensities(self, abs_normal_intensities, abs_intensities,
                           mus, pblum_scale, extinct_factors, boost_factors,
                           roche_coords, s=[0., 0., 1.], t=None):
        import numpy as np
        f = 1 + self.kwargs['amplitude'] * np.sin(2 * np.pi * (t - self.kwargs['t0']) / self.kwargs['period'])
        return abs_normal_intensities*f, abs_intensities*f

b.add_feature(SinusoidalIntensities, component='primary', period=0.2, amplitude=0.2, feature='sinusoidal_intensities')
#b.add_dataset('mesh', compute_phases=phoebe.linspace(0,1,51), columns='intensities*')
b.run_compute(model='with_sin_intens')

#b.plot(kind='mesh', animate=True, fc='intensities@lc01', ec='face', save='./mesh_with_sin_intens.gif')
b.plot(kind='lc', show=True, legend=True)

Example of a ComponentFeature (DifferentialRotation):

import phoebe
from phoebe.features import ComponentFeature
from phoebe.parameters import FloatParameter, ParameterSet
from astropy import units as u
import numpy as np

b = phoebe.default_binary()
b.add_dataset('rv', compute_phases=phoebe.linspace(0,1,101))
b.set_value('irrad_method', 'none')

b.run_compute(model='no_features')


class DifferentialRotation(ComponentFeature):
    @classmethod
    def create_feature_parameters(self, feature, **kwargs):
        params = []
        params += [FloatParameter(qualifier='alpha',
                                  latexfmt=r'\alpha_\mathrm{{ {feature} }}',
                                  value=kwargs.get('alpha', 0.2),
                                  default_unit=u.dimensionless_unscaled,
                                  description='Differential rotation coefficient')]

        return ParameterSet(params), []
    
    @classmethod
    def parse_bundle(cls, b, feature_ps):
        t0 = b.get_value(qualifier='t0', context='system', unit=u.d, **_skip_filter_checks)
        feature_dict =  cls.parse_from_feature_ps(b, feature_ps,
                                                  ['alpha'])
        return dict(t0=t0, **feature_dict)

    def requires_remeshing(self):
        return False
    
    def modify_rvs(self, rvs, orbit_vel,
                   roche_coords, s=[0., 0., 1.], t=None):
        import numpy as np

        # NOTE: orbit_vel[2] (vw) is in opposite direction of rvs!
        rvs_rot = rvs + orbit_vel[2]

        r, colat, long = self.cartesian_to_spherical(roche_coords)
        # Reiners & Royer A&A 2004
        # alpha = 0.2 for solar
        # vel_rigid(colat) = vel_equator * sin(colat)
        # vel_diffrot(colat) = vel_equator * (1-alpha sin2 (colat))
        
        sin_colat = np.sin(colat)
        eps = 1e-12
        factor = (1 - self.kwargs['alpha'] * sin_colat**2) / np.where(np.abs(sin_colat) < eps, eps, sin_colat)
        
        return -orbit_vel[2] + rvs_rot * factor

b.add_feature(DifferentialRotation, component='primary', alpha=0.8, feature='sinusoidal_intensities')
b.add_dataset('mesh', compute_phases=phoebe.linspace(0,1,51), columns=['rvs*'])
b.run_compute(model='with_diff_rot')

#b.plot(kind='mesh', animate=True, fc='rvs', ec='face', save='./mesh_with_diff_rot.gif')
b.plot(kind='rv', show=True, legend=True)

Example of a ComponentFeature inheriting from a built-in Spot (MigratingSpot):

import phoebe
from phoebe.features import Spot
from phoebe.parameters import FloatParameter, ParameterSet
from astropy import units as u
import numpy as np

b = phoebe.default_binary()
b.add_dataset('lc', compute_phases=phoebe.linspace(0,1,101))
b.set_value('irrad_method', 'none')
b.set_value_all('distortion_method', 'sphere')

b.add_feature('spot', component='primary', relteff=0.9, radius=30, colat=45, long=90, feature='static_spot')
b.run_compute(model='with_static_spot')

class MigratingSpot(Spot):
    @classmethod
    def create_feature_parameters(self, feature, **kwargs):
        from phoebe.parameters.feature import spot
        from phoebe import u

        params = []
        ps, constraints = spot(feature, **kwargs)
        params = ps.to_list()

        params += [FloatParameter(qualifier='dlongdt',
                                  value=kwargs.get('dlongdt', 0.0),
                                  default_unit=u.deg/u.d,
                                  description='Time derivative of long')]
        params += [FloatParameter(qualifier='dcolatdt',
                                  value=kwargs.get('dcolatdt', 0.0),
                                  default_unit=u.deg/u.d,
                                  description='Time derivative of colat')]

        return ParameterSet(params), constraints
    
    @classmethod
    def parse_bundle(cls, b, feature_ps):
        from phoebe.features import Spot
        spot_kw = Spot.parse_bundle(b, feature_ps)
        addl_kw = cls.parse_from_feature_ps(b, feature_ps,
                                            [{'qualifier': 'dlongdt', 'unit': 'rad/d'},
                                             {'qualifier': 'dcolatdt', 'unit': 'rad/d'}])
        return dict(**spot_kw, **addl_kw)
  
    def instantaneous_position(self, s, time):
        """
        s is the spin vector in roche coordinates
        time is the current time
        """
        t = time - self.kwargs['t0']
        longitude = self.kwargs['longitude'] + (self.kwargs['rot_dlongdt'] + self.kwargs['dlongdt']) * t
        colat = self.kwargs['colat'] + self.kwargs['dcolatdt'] * t
        return longitude, colat

b.disable_feature('static_spot')
b.add_feature(MigratingSpot, component='primary', relteff=0.9, radius=30, colat=45, long=90, dcolatdt=60, feature='migrating_spot')
b.add_dataset('mesh', compute_phases=phoebe.linspace(0,1,51), columns='intensities*')
b.run_compute(model='with_migrating_spot')

#b.plot(kind='mesh', animate=True, fc='intensities@lc01', ec='face', save='./mesh_with_migrating_spot.gif')
b.plot(kind='lc', show=True, legend=True)

TODO:

• migrate recent spot bugfixes
• support for component features (and example)
• can we extend DatasetFeature to handle modifying the comparison data (ie detrending)? Or should we split into DatasetFeature and ModelFeature?
• allow modify_intensities to have access to projected absolute intensities automatically (perhaps split into multiple hooks or at least give access to mu)
• can we be more consistent between returning values in ComponentFeature vs modifying parameters in DatasetFeature
• helper-function/access to long/colat for mesh locations
• write general LightCurveAdditiveFeature and LightCurveMultiplicativeFeature which allows just defining the functional form and handles modify_data_for_estimators and modify_model
• allow features to hook into feature checks
• replace rv_offset with feature?
• rename get_parameters to add_feature or similar?
• rename proto_coords to something more user-friendly (and document better)
• change namespace of returned class from feature_code, if possible
• access to logger and uncomment logger in spot
• update and test pulsations implementation
• update and test GPs implementation
• ensure registered internal features still show in list_available_features
• consider having built-in features in categories (component.spot, dataset.gaussian_process, etc)
• try to remove need to define imports in-method (or document otherwise)
• example for overwriting feature_code through parameter
• example for overwriting feature including parameterization (perhaps just with add_feature(..., overwrite=True, **b.get_feature(...)), or should we split into two input arguments (if callable is sent to kind then create the parameters and create a feature_code parameter which can also be sent as a kwarg).
• examples with increasing constraint-complexity (ie no constraints, built-in constraints, custom constraints)
• Remove feature registration logic (in favor of serialization)
• entries need to show up in list_available_features (at least if using registration)
• example for migrating spots
• example for sinusoidal intensities
• user-access to instantaneous properties (orbital elements, distance, anomalies, etc)
• test coverage
• documentation