Chapter 6. Workshop Tutorial, Real Example¶

# Install SciUnit if necessary
!pip install -q sciunit

# Import the package
import sciunit

# Some imports to make the code below run
from math import pi, sqrt, sin, cos, tan, atan
from datetime import datetime, timedelta
import numpy as np
# SymPy is needed because one of Kepler's equations
# is in implicit form and must be solved numerically!
from sympy import Symbol, sin as sin_
from sympy.solvers.solvers import nsolve

class ProducesOrbitalPosition(sciunit.Capability):
    """
    A model `capability`, i.e. a collection of methods that a test is allowed to invoke on a model.
    These methods are unimplemented by design, and the model must implement them.
    """
    
    def get_position(self, t: datetime) -> tuple:
        """Produce an orbital position from a time point
        in polar coordinates.
        
        Args:
            t (datetime): The time point to examine, relative to perihelion
        Returns:
            tuple: A pair of (r, theta) coordinates in the oribtal plane
        """
        raise NotImplementedError("")
        
    @property
    def perihelion(self) -> datetime:
        """Return the time of last perihelion"""
        raise NotImplementedError("")
    
    @property
    def period(self) -> float:
        """Return the period of the orbit"""
        raise NotImplementedError("")
    
    @property
    def eccentricity(self) -> float:
        """Return the eccentricity of the orbit"""
        raise NotImplementedError("")
    
    def get_x_y(self, t: datetime) -> tuple:
        """Produce an orbital position from a time point, but in cartesian coordinates.
        This method does not require a model-specific implementation.
        Thus, a generic implementation can be provided in advance."""
        raise NotImplementedError("")

# An extremely generic model capability
from sciunit.capabilities import ProducesNumber

# A specific model capability used in neurophysiology
#from neuronunit.capabilities import HasMembranePotential

class BaseKeplerModel(sciunit.Model, 
                      ProducesOrbitalPosition):
    """A sciunit model class corresponding to a Kepler-type model
    of an object in the solar system. This model has the 
    `ProducesOrbitalPosition` capability by inheritance,
    so it must implement all of the unimplemented methods of that capability"""
    
    def get_position(self, t):
        """Implementation of polar coordinate position as a function of time"""
        r, theta = self.heliocentric_distance(t), self.true_anomaly(t)
        return r, theta
    
    @property
    def perihelion(self):
        """Implementation of time of last perihelion"""
        return self.params['perihelion']
    
    @property
    def period(self):
        """Implementation of period of the orbit"""
        return self.params['period']
    
    @property
    def eccentricity(self):
        """Implementation of orbital eccentricity (assuming elliptic orbit)"""
        a, b = self.params['semimajor_axis'], self.params['semiminor_axis']
        return sqrt(1 - (b/a)**2)
    
    def get_x_y(self, t: datetime) -> tuple:
        """Produce an orbital position from a time point, but in cartesian coordinates.
        This method does not require a model-specific implementation.
        Thus, a generic implementation can be provided in advance."""
        r, theta = self.get_position(t)
        x, y = r*cos(theta), r*sin(theta)
        return x, y

class KeplerModel(BaseKeplerModel):
    """This 'full' model contains all of the methods required
    to complete the implementation of the `ProducesOrbitalPosition` capability"""
    
    def mean_anomaly(self, t):
        """How long into its period the object is at time `t`"""
        time_since_perihelion = t - self.perihelion
        return 2*pi*(time_since_perihelion % self.period)/self.period
    
    def eccentric_anomaly(self, t):
        """How far the object has gone into its period at time `t`"""
        E = Symbol('E')
        M, e = self.mean_anomaly(t), self.eccentricity
        expr = E - e*sin_(E) - M
        return nsolve(expr, 0)
    
    def true_anomaly(self, t):
        """Theta in a polar coordinate system at time `t`"""
        e, E = self.eccentricity, self.eccentric_anomaly(t)
        theta = 2*atan(sqrt(tan(E/2)**2 * (1+e)/(1-e)))
        return theta
    
    def heliocentric_distance(self, t):
        """R in a polar coordinate system at time `t`"""
        a, e = self.params['semimajor_axis'], self.eccentricity
        E = self.eccentric_anomaly(t)
        return a*(1-e*cos(E))

# The quantities module to put dimensional units on values
import quantities as pq

# `earth_model` will be a specific instance of KeplerModel, with its own parameters
earth_model = KeplerModel(name = "Kepler's Earth Model",
                          semimajor_axis=149598023 * pq.km, 
                          semiminor_axis=149577161 * pq.km, 
                          period=timedelta(365, 22118), # Period of Earth's orbit 
                          perihelion=datetime(2019, 1, 3, 0, 19), # Time and date of Earth's last perihelion
                         )

# Several score types available in SciUnit
from sciunit.scores import BooleanScore, ZScore, RatioScore, PercentScore # etc., etc.

class PositionTest(sciunit.Test):
    """A test of a planetary position at some specified time"""
    
    # This test can only operate on models that implement
    # the `ProducesOrbitalPosition` capability.
    required_capabilities = (ProducesOrbitalPosition,)
    score_type = BooleanScore # This test's 'judge' method will return a BooleanScore.
    
    def generate_prediction(self, model):
        """Generate a prediction from a model"""
        t = self.observation['t'] # Get the time point from the test's observation
        x, y = model.get_x_y(t) # Get the predicted x, y coordinates from the model
        return {'t': t, 'x': x, 'y': y} # Roll this into a model prediction dictionary
        
    def compute_score(self, observation, prediction):
        """Compute a test score based on the agreement between
        the observation (data) and prediction (model)"""
        # Compare observation and prediction to get an error measure
        delta_x = observation['x'] - prediction['x']
        delta_y = observation['y'] - prediction['y']
        error = np.sqrt(delta_x**2 + delta_y**2)
        
        passing = bool(error < 1e5*pq.kilometer) # Turn this into a True/False score
        score = self.score_type(passing) # Create a sciunit.Score object
        score.set_raw(error) # Add some information about how this score was obtained
        score.description = ("Passing score if the prediction is "
                             "within < 100,000 km of the observation") # Describe the scoring logic
        return score

class StricterPositionTest(PositionTest):
    # Optional observation units to validate against
    units = pq.meter
    
    # Optional schema for the format of observed data
    observation_schema = {'t': {'min': 0, 'required': True},
                          'x': {'units': True, 'required': True},
                          'y': {'units': True, 'required': True},
                          'phi': {'required': False}}
    
    def validate_observation(self, observation):
        """Additional checks on the observation"""
        assert isinstance(observation['t'], datetime)
        return observation
        
    # Optional schema for the format of test parameters
    params_schema = {'rotate': {'required': False}}

    # Optional schema for the format of default test parameters
    default_params = {'rotate': False}
    
    def compute_score(self, observation, prediction):
        """Optionally use additional information to compute model/data agreement"""
        observation_rotated = observation.copy()
        if 'phi' in observation:
            # Project x and y values onto the plane defined by `phi`.
            observation_rotated['x'] *= cos(observation['phi'])
            observation_rotated['y'] *= cos(observation['phi'])
        return super().compute_score(observation_rotated, prediction)
            

# A single test instance, best on the test class `StricterPositionTest` combined with
# a specific set of observed data (a time and some x, y coordinates)
# N.B.: This data is made up for illustration purposes
earth_position_test_march = StricterPositionTest(name = "Earth Orbital Data on March 1st, 2019",
                                                 observation = {'t': datetime(2019, 3, 1),
                                                                'x': 7.905e7 * pq.km, 
                                                                'y': 1.254e8 * pq.km})

# A new test for a new month: same test class, new observation (data)
# N.B. I deliberately picked "observed" values that will make the model fail this test
earth_position_test_april = StricterPositionTest(name = "Earth Orbital Data on April 1st, 2019",
                                                 observation = {'t': datetime(2019, 4, 1),
                                                                'x': 160000 * pq.km, 
                                                                'y': 70000 * pq.km})

# A test suite built from both of the tests that we have instantiated
earth_position_suite = sciunit.TestSuite([earth_position_test_march,
                                          earth_position_test_april],
                                         name = 'Earth observations in Spring, 2019')

# Just like the Kepler model, but returning a random orbital angle
class RandomModel(KeplerModel):
    def get_position(self, t):
        r, theta = super().get_position(t)
        return r, 2*pi*np.random.rand()

# A new model instance, using the same parameters but a different underlying model class
random_model = RandomModel(name = "Random Earth Model",
                           semimajor_axis=149598023 * pq.km, 
                           semiminor_axis=149577161 * pq.km, 
                           period=timedelta(365, 22118), # Period of Earth's orbit 
                           perihelion=datetime(2019, 1, 3, 0, 19), # Time and date of Earth's last perihelion
                           )

# A simple model which has some capabilities, 
# but not the ones needed for the orbital position test
class SimpleModel(sciunit.Model,
                  sciunit.capabilities.ProducesNumber):
    pass

simple_model = SimpleModel()