<commit_before>"""Attitude reconstruction of an ArduPilot Mega."""


import os
import re

import yaipopt
import numpy as np
import sympy
import sym2num
from numpy import ma
from scipy import io

from ceacoest import kalman, sde, utils


class SymbolicModel(sde.SymbolicModel):
    """Symbolic SDE model."""
    
    var_names = {'t', 'x', 'y', 'q', 'c'}
    """Name of model variables."""
    
    function_names = {'f', 'g', 'h', 'R'}
    """Name of the model functions."""
    
    t = 't'
    """Time variable."""
    
    x = ['ax', 'ay', 'az', 'p', 'q', 'r', 'q0', 'q1', 'q2', 'q3']
    """State vector."""
    
    y = ['ax_meas', 'ay_meas', 'az_meas', 'p_meas', 'q_meas', 'r_meas',
         'magx_meas', 'magy_meas', 'magz_meas']
    """Measurement vector."""
    
    q = ['acc_png', 'omega_png', 'acc_mng', 'omega_mng', 'mag_mng',
         'magex', 'magez']
    """Parameter vector."""
    
    c = ['quat_renorm_gain', 'g0']
    """Constants vector."""
    
    def f(self, t, x, q, c):
        """Drift function."""
        s = self.symbols(t=t, x=x, q=q, c=c)
        renorm = s.quat_renorm_gain*(1 - s.q0**2 - s.q1**2 - s.q2**2 - s.q3**2)
        derivs = dict(
            p=0, q=0, r=0, ax=0, ay=0, az=0,
            q0=-0.5*(s.q1*s.p + s.q2*s.q + s.q3*s.r) + renorm*s.q0,
            q1=-0.5*(-s.q0*s.p - s.q2*s.r + s.q3*s.q) + renorm*s.q1,
            q2=-0.5*(-s.q0*s.q + s.q1*s.r - s.q3*s.p) + renorm*s.q2,
            q3=-0.5*(-s.q0*s.r - s.q1*s.q + s.q2*s.p) + renorm*s.q3,
        )
        return self.pack('x', derivs)
    
    def g(self, t, x, q, c):
        """Diffusion matrix."""
        s = self.symbols(t=t, x=x, q=q, c=c)
        g = np.zeros((x.size, 6), object)
        g[[0, 1, 2], [0, 1, 2]] = s.acc_png
        g[[3, 4, 5], [3, 4, 5]] = s.omega_png
        return g
    
    def h(self, t, x, q, c):
        """Measurement function."""
        s = self.symbols(t=t, x=x, q=q, c=c)
        q0 = s.q0
        q1 = s.q1
        q2 = s.q2
        q3 = s.q3
        e2b = np.array(
            [[q0**2+q1**2-q2**2-q3**2, 2*(q1*q2 + q0*q3), 2*(q1*q3 - q0*q2)],
             [2*(q1*q2 - q0*q3), q0**2-q1**2+q2**2-q3**2, 2*(q2*q3 + q0*q1)],
             [2*(q1*q3 + q0*q2), 2*(q2*q3 - q0*q1), q0**2-q1**2-q2**2+q3**2]],
            dtype=object
        )
        magb = np.dot(e2b, [s.magex, 0, s.magez])
        gb = np.dot(e2b, [0, 0, s.g0])
        meas = dict(
            p_meas=s.p, q_meas=s.q, r_meas=s.r,
            magx_meas=magb[0], magy_meas=magb[1], magz_meas=magb[2],
            ax_meas=s.ax - gb[0], 
            ay_meas=s.ay - gb[1], 
            az_meas=s.az - gb[2],
        )
        return self.pack('y', meas)
    
    def R(self, q, c):
        """Measurement function."""
        s = self.symbols(q=q, c=c)
        std = dict(
            ax_meas=s.acc_mng, ay_meas=s.acc_mng, az_meas=s.acc_mng, 
            p_meas=s.omega_mng, q_meas=s.omega_mng, r_meas=s.omega_mng,
            magx_meas=s.mag_mng, magy_meas=s.mag_mng, magz_meas=s.mag_mng,
        )
        return np.diag(self.pack('y', std) ** 2)


class SymbolicDTModel(SymbolicModel, sde.ItoTaylorAS15DiscretizedModel):
    derivatives = [('df_dx', 'f', 'x'), ('df_dq', 'f', 'q'),
                   ('d2f_dx2', 'df_dx',  'x'), 
                   ('d2f_dx_dq', 'df_dx', 'q'),
                   ('d2f_dq2', 'df_dq',  'q'),
                   ('dQ_dx', 'Q', 'x'), ('dQ_dq', 'Q', 'q'),
                   ('d2Q_dx2', 'dQ_dx',  'x'), 
                   ('d2Q_dx_dq', 'dQ_dx', 'q'),
                   ('d2Q_dq2', 'dQ_dq',  'q'),
                   ('dh_dx', 'h', 'x'), ('dh_dq', 'h', 'q'),
                   ('d2h_dx2', 'dh_dx',  'x'), 
                   ('d2h_dx_dq', 'dh_dx', 'q'),
                   ('d2h_dq2', 'dh_dq',  'q'),
                   ('dR_dq', 'R', 'q'), ('d2R_dq2', 'dR_dq', 'q')]
    """List of the model function derivatives to calculate / generate."""
    
    dt = 'dt'
    """Discretization time step."""
    
    k = 'k'
    """Discretized sample index."""
    
    generated_name = "GeneratedDTModel"
    """Name of the generated class."""
    
    meta = 'ceacoest.sde.DiscretizedModel.meta'
    """Generated model metaclass."""
    
    @property
    def imports(self):
        return super().imports + ('import ceacoest.sde',)


sym_model = SymbolicModel()
sym_dt_model = SymbolicDTModel()
printer = sym2num.ScipyPrinter()
GeneratedDTModel = sym2num.class_obj(sym_dt_model, printer)


def load_data():
    module_dir = os.path.dirname(__file__)
    filepath = os.path.join(module_dir, 'data', 'apm.log')
    lines = open(filepath).read().splitlines()
    
    t = []
    y = []
    for line in lines:
        msgid, *fields = re.split(',\s*', line)
        if msgid == 'MAG':
            t.append(float(fields[0]) / 1e3)
            ydict = dict(
                magx_meas=float(fields[1]),
                magy_meas=float(fields[2]),
                magz_meas=float(fields[3]),
            )
            y.append(GeneratedDTModel.pack('y', ydict, fill=np.nan))
        elif msgid == 'IMU':
            t.append(float(fields[0]) / 1e3)
            ydict = dict(
                p_meas=float(fields[1]),
                q_meas=float(fields[2]),
                r_meas=float(fields[3]),
                ax_meas=float(fields[4]),
                ay_meas=float(fields[5]),
                az_meas=float(fields[6]),
            )
            y.append(GeneratedDTModel.pack('y', ydict, fill=np.nan))
    range_ = np.s_[900:1000]#np.s_[900:1800]
    t = np.asarray(t)[range_]
    y = ma.masked_invalid(y)[range_]
    assert np.unique(t).size == t.size
    return t, y


def pem(t, y):
    # Instantiate the model
    given = dict(
        g0=9.81, quat_renorm_gain=4,
        acc_png=0.01, omega_png=0.0001,
        acc_mng=0.03, omega_mng=8e-3, mag_mng=1.6,
        magex=117, magez=0,
    )
    dt = np.diff(t)
    q0 = GeneratedDTModel.pack('q', given)
    c = GeneratedDTModel.pack('c', given)
    params = dict(q=q0, c=c)
    sampled = dict(dt=dt, t=t)
    model = GeneratedDTModel(params, sampled)

    # Build the initial state and covariance
    x0 = np.zeros(GeneratedDTModel.nx)
    x0[-4] = 1
    Px0 = np.diag(np.repeat([1, 1e-3, 0.1], (3,3,4)) ** 2)

    def merit(q, new=None):
        mq = model.parametrize(q=q)
        kf = kalman.DTUnscentedKalmanFilter(mq, x0, Px0)
        return kf.pem_merit(y)
    
    def grad(q, new=None):
        mq = model.parametrize(q=q)
        kf = kalman.DTUnscentedKalmanFilter(mq, x0, Px0)
        return kf.pem_gradient(y)
    
    hess_inds = np.tril_indices(model.nq)
    def hess(q, new_q=1, obj_factor=1, lmult=1, new_lmult=1):
        mq = model.parametrize(q=q)
        kf = kalman.DTUnscentedKalmanFilter(mq, x0, Px0)
        return obj_factor * kf.pem_hessian(y)[hess_inds]
    
    q_lb = dict(ang_meas_std=0, angvel_png=0)
    q_ub = dict()
    q_fix = dict(acc_mng=0.03, omega_mng=8e-3, mag_mng=1.6)
    q_bounds = [model.pack('q', dict(q_lb, **q_fix), fill=-np.inf),
                model.pack('q', dict(q_ub, **q_fix), fill=np.inf)]
    problem = yaipopt.Problem(q_bounds, merit, grad,
                            hess=hess, hess_inds=hess_inds)
    problem.num_option(b'obj_scaling_factor', -1)
    (qopt, solinfo) = problem.solve(q0)
    
    return problem, qopt, solinfo, model, q0, x0, Px0


if __name__ == '__main__':
    [t, y] = load_data()
    
    
<commit_msg>examples: added hard iron to apm_attitude.<commit_after>"""Attitude reconstruction of an ArduPilot Mega."""


import os
import re

import yaipopt
import numpy as np
import sympy
import sym2num
from numpy import ma

from ceacoest import kalman, sde, utils


class SymbolicModel(sde.SymbolicModel):
    """Symbolic SDE model."""
    
    var_names = {'t', 'x', 'y', 'q', 'c'}
    """Name of model variables."""
    
    function_names = {'f', 'g', 'h', 'R'}
    """Name of the model functions."""
    
    t = 't'
    """Time variable."""
    
    x = ['ax', 'ay', 'az', 'p', 'q', 'r', 'q0', 'q1', 'q2', 'q3']
    """State vector."""
    
    y = ['ax_meas', 'ay_meas', 'az_meas', 'p_meas', 'q_meas', 'r_meas',
         'magx_meas', 'magy_meas', 'magz_meas']
    """Measurement vector."""
    
    q = ['acc_png', 'omega_png', 'acc_mng', 'omega_mng', 'mag_mng',
         'magex', 'magez', 'maghx', 'maghy', 'maghz']
    """Parameter vector."""
    
    c = ['quat_renorm_gain', 'g0']
    """Constants vector."""
    
    def f(self, t, x, q, c):
        """Drift function."""
        s = self.symbols(t=t, x=x, q=q, c=c)
        renorm = s.quat_renorm_gain*(1 - s.q0**2 - s.q1**2 - s.q2**2 - s.q3**2)
        derivs = dict(
            p=0, q=0, r=0, ax=0, ay=0, az=0,
            q0=-0.5*(s.q1*s.p + s.q2*s.q + s.q3*s.r) + renorm*s.q0,
            q1=-0.5*(-s.q0*s.p - s.q2*s.r + s.q3*s.q) + renorm*s.q1,
            q2=-0.5*(-s.q0*s.q + s.q1*s.r - s.q3*s.p) + renorm*s.q2,
            q3=-0.5*(-s.q0*s.r - s.q1*s.q + s.q2*s.p) + renorm*s.q3,
        )
        return self.pack('x', derivs)
    
    def g(self, t, x, q, c):
        """Diffusion matrix."""
        s = self.symbols(t=t, x=x, q=q, c=c)
        g = np.zeros((x.size, 6), object)
        g[[0, 1, 2], [0, 1, 2]] = s.acc_png
        g[[3, 4, 5], [3, 4, 5]] = s.omega_png
        return g
    
    def h(self, t, x, q, c):
        """Measurement function."""
        s = self.symbols(t=t, x=x, q=q, c=c)
        q0 = s.q0
        q1 = s.q1
        q2 = s.q2
        q3 = s.q3
        e2b = np.array(
            [[q0**2+q1**2-q2**2-q3**2, 2*(q1*q2 + q0*q3), 2*(q1*q3 - q0*q2)],
             [2*(q1*q2 - q0*q3), q0**2-q1**2+q2**2-q3**2, 2*(q2*q3 + q0*q1)],
             [2*(q1*q3 + q0*q2), 2*(q2*q3 - q0*q1), q0**2-q1**2-q2**2+q3**2]],
            dtype=object
        )
        mag = np.dot(e2b, [s.magex, 0, s.magez]) + [s.maghx, s.maghy, s.maghz]
        gb = np.dot(e2b, [0, 0, s.g0])
        meas = dict(
            p_meas=s.p, q_meas=s.q, r_meas=s.r,
            magx_meas=mag[0], magy_meas=mag[1], magz_meas=mag[2],
            ax_meas=s.ax - gb[0],
            ay_meas=s.ay - gb[1],
            az_meas=s.az - gb[2],
        )
        return self.pack('y', meas)
    
    def R(self, q, c):
        """Measurement function."""
        s = self.symbols(q=q, c=c)
        std = dict(
            ax_meas=s.acc_mng, ay_meas=s.acc_mng, az_meas=s.acc_mng, 
            p_meas=s.omega_mng, q_meas=s.omega_mng, r_meas=s.omega_mng,
            magx_meas=s.mag_mng, magy_meas=s.mag_mng, magz_meas=s.mag_mng,
        )
        return np.diag(self.pack('y', std) ** 2)


class SymbolicDTModel(SymbolicModel, sde.ItoTaylorAS15DiscretizedModel):
    derivatives = [('df_dx', 'f', 'x'), ('df_dq', 'f', 'q'),
                   ('d2f_dx2', 'df_dx',  'x'), 
                   ('d2f_dx_dq', 'df_dx', 'q'),
                   ('d2f_dq2', 'df_dq',  'q'),
                   ('dQ_dx', 'Q', 'x'), ('dQ_dq', 'Q', 'q'),
                   ('d2Q_dx2', 'dQ_dx',  'x'), 
                   ('d2Q_dx_dq', 'dQ_dx', 'q'),
                   ('d2Q_dq2', 'dQ_dq',  'q'),
                   ('dh_dx', 'h', 'x'), ('dh_dq', 'h', 'q'),
                   ('d2h_dx2', 'dh_dx',  'x'), 
                   ('d2h_dx_dq', 'dh_dx', 'q'),
                   ('d2h_dq2', 'dh_dq',  'q'),
                   ('dR_dq', 'R', 'q'), ('d2R_dq2', 'dR_dq', 'q')]
    """List of the model function derivatives to calculate / generate."""
    
    dt = 'dt'
    """Discretization time step."""
    
    k = 'k'
    """Discretized sample index."""
    
    generated_name = "GeneratedDTModel"
    """Name of the generated class."""
    
    meta = 'ceacoest.sde.DiscretizedModel.meta'
    """Generated model metaclass."""
    
    @property
    def imports(self):
        return super().imports + ('import ceacoest.sde',)


sym_model = SymbolicModel()
sym_dt_model = SymbolicDTModel()
printer = sym2num.ScipyPrinter()
GeneratedDTModel = sym2num.class_obj(sym_dt_model, printer)


def load_data():
    module_dir = os.path.dirname(__file__)
    filepath = os.path.join(module_dir, 'data', 'apm.log')
    lines = open(filepath).read().splitlines()
    
    t = []
    y = []
    for line in lines:
        msgid, *fields = re.split(',\s*', line)
        if msgid == 'MAG':
            t.append(float(fields[0]) / 1e3)
            ydict = dict(
                magx_meas=float(fields[1]),
                magy_meas=float(fields[2]),
                magz_meas=float(fields[3]),
            )
            y.append(GeneratedDTModel.pack('y', ydict, fill=np.nan))
        elif msgid == 'IMU':
            t.append(float(fields[0]) / 1e3)
            ydict = dict(
                p_meas=float(fields[1]),
                q_meas=float(fields[2]),
                r_meas=float(fields[3]),
                ax_meas=float(fields[4]),
                ay_meas=float(fields[5]),
                az_meas=float(fields[6]),
            )
            y.append(GeneratedDTModel.pack('y', ydict, fill=np.nan))
    range_ = np.s_[900:1200]#np.s_[900:1800]
    t = np.asarray(t)[range_]
    y = ma.masked_invalid(y)[range_]
    assert np.unique(t).size == t.size
    return t, y


def pem(t, y):
    # Instantiate the model
    given = dict(
        g0=9.81, quat_renorm_gain=4,
        acc_png=0.01, omega_png=0.0001,
        acc_mng=0.03, omega_mng=8e-3, mag_mng=1.6,
        magex=117, magez=0, maghx=0, maghy=0, maghz=0,
    )
    dt = np.diff(t)
    q0 = GeneratedDTModel.pack('q', given)
    c = GeneratedDTModel.pack('c', given)
    params = dict(q=q0, c=c)
    sampled = dict(dt=dt, t=t)
    model = GeneratedDTModel(params, sampled)

    # Build the initial state and covariance
    x0 = np.zeros(GeneratedDTModel.nx)
    x0[-4] = 1
    Px0 = np.diag(np.repeat([1, 1e-3, 0.1], (3,3,4)) ** 2)

    def merit(q, new=None):
        mq = model.parametrize(q=q)
        kf = kalman.DTUnscentedKalmanFilter(mq, x0, Px0)
        return kf.pem_merit(y)
    
    def grad(q, new=None):
        mq = model.parametrize(q=q)
        kf = kalman.DTUnscentedKalmanFilter(mq, x0, Px0)
        return kf.pem_gradient(y)
    
    hess_inds = np.tril_indices(model.nq)
    def hess(q, new_q=1, obj_factor=1, lmult=1, new_lmult=1):
        mq = model.parametrize(q=q)
        kf = kalman.DTUnscentedKalmanFilter(mq, x0, Px0)
        return obj_factor * kf.pem_hessian(y)[hess_inds]
    
    q_lb = dict(ang_meas_std=0, angvel_png=0)
    q_ub = dict()
    q_fix = dict(acc_mng=0.03, omega_mng=8e-3, mag_mng=1.6)
    q_bounds = [model.pack('q', dict(q_lb, **q_fix), fill=-np.inf),
                model.pack('q', dict(q_ub, **q_fix), fill=np.inf)]
    problem = yaipopt.Problem(q_bounds, merit, grad,
                            hess=hess, hess_inds=hess_inds)
    problem.num_option(b'obj_scaling_factor', -1)
    (qopt, solinfo) = problem.solve(q0)
    
    return problem, qopt, solinfo, model, q0, x0, Px0


if __name__ == '__main__':
    [t, y] = load_data()
    
    

{'fraction_non_alphanumeric': 0.18897584671434353, 'fraction_numerical': 0.03420591013358521, 'mean_word_length': 2.0727611940298507, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6403730', 'n_tokens_mistral': 6491, 'n_tokens_neox': 6153, 'n_words': 1144}