def th_dihedral_angle(r0,r1,r2,r3, shift=0.):
    bxa = thcross(r2 - r1, r1 - r0)
    bxa /= bxa.norm(2)
    cxb = thcross(r3 - r2, r2 - r1)
    cxb /= cxb.norm(2)
    angle = th.arccos(th.dot(bxa, cxb))
    return angle + shift
Esempio n. 2
0
    def loss(lv1, lv2):
        """ Contrastive cosine distance optimization target """

        # number of samples in batch
        n = lv1.shape[0]

        # compute cosine distance
        D = lv1.dot(lv2.T)

        # compute arcus cosinus -> converts similarity into distance
        D = T.arccos(D)

        # distance between matching pairs
        d = D.diagonal().reshape((-1, 1))

        # distance between non-matching pairs
        M = T.identity_like(D)
        O = D[(M <= 0).nonzero()].reshape((n, n - 1))

        # max margin hinge loss
        L = gamma + d
        L = T.repeat(L, n - 1, 1)
        L -= O
        L = T.clip(L, 0, 1000)

        # compute batch mean
        loss = L.mean()

        return weight * loss
def th_dihedral_angle_rec(r0,r1,r2,r3, shift=0.):
    # reciprocal of the angle
    bxa = thcross(r2 - r1, r1 - r0)
    bxa /= bxa.norm(2)
    cxb = thcross(r3 - r2, r2 - r1)
    cxb /= cxb.norm(2)
    angle = 2*sp.pi - th.arccos(th.dot(bxa, cxb))
    return angle + shift
 def generator(self, u):
     if u.ndim == 1:
         u = u[None, :]
     z, n = u[:, :self.z_dim], u[:, self.z_dim:self.z_dim + self.x_dim]
     loc, scale = self.x_gvn_z(z)
     alpha = 2 * scale / (1 + scale**2)
     phi = tt.arccos((tt.cos(2 * np.pi * n) + alpha) /
                     (1 + alpha * tt.cos(2 * np.pi * n))) + loc
     return tt.switch(n < 0.5, phi, 2 * np.pi - phi)
Esempio n. 5
0
    def log(self, X, Y):
        V = self.proj(X, Y - X)
        dists = tensor.arccos((X * Y).sum(0))
        norms = tensor.sqrt((V ** 2).sum(0)).real
        factors = dists / norms
        # For very close points, dists is almost equal to norms, but because
        # they are both almost zero, the division above can return NaN's. To
        # avoid that, we force those ratios to 1.
        factors[dists <= 1e-6] = 1

        return V * factors
Esempio n. 6
0
 def loss(self, y, output):
     print 'Calling SentenceLSTMLayers.loss()...'
     # y is a single output. No support for minibatches as of yet.
     W_dist, W_angle = 0.5, 0.5
     assert output.ndim == 1
     y_ = assert_op(y, T.eq(y.shape[0], output.shape[0]), T.eq(y.shape[1], 1))
     loss = T.mean((output - y_[:, 0])**2) + T.arccos(T.dot(output, y[:, 0])/(absolute(output)*absolute(y[:, 0])))
     loss_ = assert_op(loss, T.eq(loss.ndim, 0))
     print 'The dimensions of loss is', loss.ndim
     print 'SentenceLSTMLayers.loss() call successful'
     return loss_
Esempio n. 7
0
 def calctheta(self, vec):
     try:
         return self.CALCTHETA(vec)
     except:  # Else build function
         VEC = T.vector()
         calctheta = T.arccos(
             T.clip(VEC[2] / T.sqrt(VEC[0]**2 + VEC[1]**2 + VEC[2]**2),
                    -1.0, 1.0))
         self.CALCTHETA = theano.function(
             [VEC], calctheta)  # function is compiled to calculate theta
         return self.CALCTHETA(vec)
Esempio n. 8
0
    def localdepth(self, theta, phi):
        try:  # If depth function has been compiled
            return self.DEPTH(theta, phi)
        except:
            ############################################################################
            ########################## Define depth function ###########################
            ############################################################################
            print "Setting up and compiling intrinsic surface function... sit tight.\n"

            self.depth = (
                (self.MAG * T.exp(-self.BETA * T.arccos(
                    T.clip(T.dot(self.Si, [self.refX, self.refY, self.refZ]),
                           -1.0, 1.0))**2)).sum() /
                (T.exp(-self.BETA * T.arccos(
                    T.clip(T.dot(self.Si, [self.refX, self.refY, self.refZ]),
                           -1.0, 1.0))**2)).sum())

            self.DEPTH = theano.function(
                [self.THETA, self.PHI],
                self.depth)  # depth function is compiled here
            return self.DEPTH(theta, phi)
    def get_arc_cosine_penalty(self):
        """Calculate the arccosine distance in [0, 1].
        0: the two vectors are very similar. (have the same orientation)
        1: the two vectors are very disimilar (have the opposite orientation).
        The cosine similarity does not take in consideration the magnitude
        of the vectors. It considers only thier orientation (angle).
        Therefore, two vectors are similar if they have the same angle.
        See: https://en.wikipedia.org/wiki/Cosine_similarity
        """
        # tiny value:
        #        tiny = sharedX_value(np.finfo(dtype=theano.config.floatX).tiny,
        #                             name="tiny")
        # the gradient of sqrt at 0 is undefinded (nan).
        # use a tiny value instead of 0.
        # OLD SOLUTION
        #        denom = T.sqrt(
        #            T.sum(self.output1**2, axis=1) * T.sum(self.output2**2, axis=1))
        #        nomin = (self.output1 * self.output2).sum(1)
        #        cosine = nomin/denom  # the cosine betwen the two vectors
        #        pi = sharedX_value(np.pi, name="pi")
        #        minus1 = sharedX_value(-1., name="minus1")
        #        plus1 = sharedX_value(1. - np.finfo(dtype=theano.config.floatX).eps,
        #                              name="plus1")
        #        # Need to be clipped. accos() gives nan when sin is close to 1.
        #        angle = T.arccos(T.clip(
        #            cosine, minus1.get_value(), plus1.get_value()))/pi
        # OLD SOLUTION
        #        plus1 = sharedX_value(1. - np.finfo(dtype=theano.config.floatX).eps,
        #                              name="plus1")
        pi = sharedX_value(np.pi, name="pi")
        cosine = T.clip(
            ((self.output1 * self.output2).sum(axis=1) /
             (self._magnitude(self.output1) * self._magnitude(self.output2))),
            -1, 1 - 1e-7)
        angle = T.clip(T.arccos(cosine) / pi, 0, 1)

        return angle
Esempio n. 10
0
def get_model(parallax=None):
    with pm.Model() as model:

        if parallax is None:
            # Without an actual parallax measurement, we can model the orbit in units of arcseconds
            # by providing a fake_parallax and conversion constant
            plx = 1  # arcsec
        else:
            # Below we will run a version of this model where a measurement of parallax is provided
            # The measurement is in milliarcsec
            m_plx = pm.Bound(pm.Normal, lower=0, upper=100)(
                "m_plx", mu=parallax[0], sd=parallax[1], testval=parallax[0]
            )
            plx = pm.Deterministic("plx", 1e-3 * m_plx)

        a_ang = pm.Uniform("a_ang", 0.1, 1.0, testval=0.324)
        a = pm.Deterministic("a", a_ang / plx)

        # We expect the period to be somewhere in the range of 25 years,
        # so we'll set a broad prior on logP
        logP = pm.Normal(
            "logP", mu=np.log(25 * yr), sd=10.0, testval=np.log(28.8 * yr)
        )
        P = pm.Deterministic("P", tt.exp(logP))

        # For astrometric-only fits, it's generally better to fit in
        # p = (Omega + omega)/2 and m = (Omega - omega)/2 instead of omega and Omega
        # directly
        omega0 = 251.6 * deg - np.pi
        Omega0 = 159.6 * deg
        p = Angle("p", testval=0.5 * (Omega0 + omega0))
        m = Angle("m", testval=0.5 * (Omega0 - omega0))
        omega = pm.Deterministic("omega", p - m)
        Omega = pm.Deterministic("Omega", p + m)

        # For these orbits, it can also be better to fit for a phase angle
        # (relative to a reference time) instead of the time of periasteron
        # passage directly
        phase = Angle("phase", testval=0.0)
        tperi = pm.Deterministic("tperi", T0 + P * phase / (2 * np.pi))

        # Geometric uiform prior on cos(incl)
        cos_incl = pm.Uniform(
            "cos_incl", lower=-1, upper=1, testval=np.cos(96.0 * deg)
        )
        incl = pm.Deterministic("incl", tt.arccos(cos_incl))
        ecc = pm.Uniform("ecc", lower=0.0, upper=1.0, testval=0.798)

        # Set up the orbit
        orbit = xo.orbits.KeplerianOrbit(
            a=a * au_to_R_sun,
            t_periastron=tperi,
            period=P,
            incl=incl,
            ecc=ecc,
            omega=omega,
            Omega=Omega,
        )
        if parallax is not None:
            pm.Deterministic("M_tot", orbit.m_total)

        # Compute the model in rho and theta
        rho_model, theta_model = orbit.get_relative_angles(astro_jds, plx)
        pm.Deterministic("rho_model", rho_model)
        pm.Deterministic("theta_model", theta_model)

        # Add jitter terms to both separation and position angle
        log_rho_s = pm.Normal(
            "log_rho_s", mu=np.log(np.median(rho_err)), sd=5.0
        )
        log_theta_s = pm.Normal(
            "log_theta_s", mu=np.log(np.median(theta_err)), sd=5.0
        )
        rho_tot_err = tt.sqrt(rho_err ** 2 + tt.exp(2 * log_rho_s))
        theta_tot_err = tt.sqrt(theta_err ** 2 + tt.exp(2 * log_theta_s))

        # define the likelihood function, e.g., a Gaussian on both rho and theta
        pm.Normal("rho_obs", mu=rho_model, sd=rho_tot_err, observed=rho_data)

        # We want to be cognizant of the fact that theta wraps so the following is equivalent to
        # pm.Normal("obs_theta", mu=theta_model, observed=theta_data, sd=theta_tot_err)
        # but takes into account the wrapping. Thanks to Rob de Rosa for the tip.
        theta_diff = tt.arctan2(
            tt.sin(theta_model - theta_data), tt.cos(theta_model - theta_data)
        )
        pm.Normal("theta_obs", mu=theta_diff, sd=theta_tot_err, observed=0.0)

        # Set up predicted orbits for later plotting
        rho_dense, theta_dense = orbit.get_relative_angles(t_fine, plx)
        rho_save = pm.Deterministic("rho_save", rho_dense)
        theta_save = pm.Deterministic("theta_save", theta_dense)

        # Optimize to find the initial parameters
        map_soln = model.test_point
        map_soln = xo.optimize(map_soln, vars=[log_rho_s, log_theta_s])
        map_soln = xo.optimize(map_soln, vars=[phase])
        map_soln = xo.optimize(map_soln, vars=[p, m, ecc])
        map_soln = xo.optimize(map_soln, vars=[logP, a_ang, phase])
        map_soln = xo.optimize(map_soln)

    return model, map_soln
Esempio n. 11
0
import numpy as np 
import scipy.signal as ss
import argparse
import theano
import theano.tensor as T

# Theano functions
# QEA function
x0, x1, x2 = (T.zvector('x0'),T.zvector('x1'), T.zvector('x2'))
IF_FUN = T.arccos((x2 + x0) / (2 * x1 + 1e-16)) / np.pi / 2
QEA_FUN = theano.function([x0, x1, x2], IF_FUN)
        

def floatx(np_data):
    return np.asarray(np_data, dtype=theano.config.floatX)


def qea(im):
    """
    Quasi-eigen approximation function
    Input
    - im : 1d vector that contains a time series
    Ouput
    - ia : instantaneous amplitude
    - ip : instantaneous phase
    - ifeq: instantaneous frequency
    """
    im = floatx(im.ravel())
    # computes analytic signal
    H = ss.hilbert(im)
    H = im+1j*H
Esempio n. 12
0
    def build_model(self):
        print('Building the model')
        self.pm_model = pm.Model()

        with self.pm_model:
            # Mode locations
            numax =  pm.Normal('numax', mu = self.init['numax'], sigma = 10., testval = self.init['numax'])
            alpha =  pm.Lognormal('alpha', mu = np.log(self.init['alpha']), sigma = 0.01, testval = self.init['alpha'])
            epsilon = pm.Normal('epsilon', mu = self.init['epsilon'], sigma = 1., testval = self.init['epsilon'])
            d01     = pm.Lognormal('d01', mu = np.log(self.init['d01']), sigma = 0.1, testval = self.init['d01'])
            d02     = pm.Lognormal('d02', mu = np.log(self.init['d02']), sigma = 0.1, testval = self.init['d02'])

            sigma0 = pm.HalfCauchy('sigma0', beta = 2., testval = self.init['sigma0'])
            sigma1 = pm.HalfCauchy('sigma1', beta = 2., testval = self.init['sigma1'])
            sigma2 = pm.HalfCauchy('sigma2', beta = 2., testval = self.init['sigma2'])

            f0 = pm.Normal('f0', mu = self.mod.f0([numax, alpha, epsilon, d01, d02]), sigma = sigma0, shape=len(self.mod.n0))
            f1 = pm.Normal('f1', mu = self.mod.f1([numax, alpha, epsilon, d01, d02]), sigma = sigma1, shape=len(self.mod.n1))
            f2 = pm.Normal('f2', mu = self.mod.f2([numax, alpha, epsilon, d01, d02]), sigma = sigma2, shape=len(self.mod.n2))

            # Mode Linewidths
            m = pm.Normal('m', self.init['m'], 1., testval = self.init['m'])
            c = pm.Normal('c', self.init['c'], 1., testval = self.init['c'])
            rho = pm.Lognormal('rho', mu = np.log(self.init['rho']), sigma = 0.1, testval = self.init['rho'])
            # ls = pm.TruncatedNormal('ls', mu = np.log(self.init['L']), sigma = 0.1, lower=0., testval = self.init['L'])

            mu = pm.gp.mean.Linear(coeffs = m, intercept = c)
            cov = tt.sqr(rho) * pm.gp.cov.ExpQuad(1, ls = self.init['L'])

            gp = pm.gp.Latent(cov_func = cov, mean_func = mu)
            lng = gp.prior('lng', X = self.nf_)

            g0 = pm.Deterministic('g0', tt.exp(lng)[0:len(f0_)])
            g1 = pm.Deterministic('g1', tt.exp(lng)[len(f0_):len(f0_)+len(f1_)])
            g2 = pm.Deterministic('g2', tt.exp(lng)[len(f0_)+len(f1_):])

            # Mode Amplitude & Height
            w = pm.Lognormal('w', mu = np.log(self.init['w']), sigma = 10., testval=self.init['w'])
            A = pm.Lognormal('A', mu = np.log(self.init['A']), sigma = 1., testval=self.init['A'])
            V1 = pm.Lognormal('V1', mu = np.log(self.init['V1']), sigma = 0.1, testval=self.init['V1'])
            V2 = pm.Lognormal('V2', mu = np.log(self.init['V2']), sigma = 0.1, testval=self.init['V2'])

            sigmaA = pm.HalfCauchy('sigmaA', beta = 1., testval = self.init['sigmaA'])
            Da0 = pm.Normal('Da0', mu = 0., sigma = 1., shape=len(self.mod.n0))
            Da1 = pm.Normal('Da1', mu = 0., sigma = 1., shape=len(self.mod.n1))
            Da2 = pm.Normal('Da2', mu = 0., sigma = 1., shape=len(self.mod.n2))

            a0 = pm.Deterministic('a0', sigmaA * Da0 + mod.A0(f0_, [numax, w, A, V1, V2]))
            a1 = pm.Deterministic('a1', sigmaA * Da1 + mod.A1(f1_, [numax, w, A, V1, V2]))
            a2 = pm.Deterministic('a2', sigmaA * Da2 + mod.A2(f2_, [numax, w, A, V1, V2]))

            h0 = pm.Deterministic('h0', 2*tt.sqr(a0)/np.pi/g0)
            h1 = pm.Deterministic('h1', 2*tt.sqr(a1)/np.pi/g1)
            h2 = pm.Deterministic('h2', 2*tt.sqr(a2)/np.pi/g2)

            # Mode splitting
            xsplit = pm.Lognormal('xsplit', mu = np.log(self.init['xsplit']), sigma=  0.75, testval = self.init['xsplit'])
            cosi = pm.Uniform('cosi', 0., 1., testval = self.init['cosi'])

            i = pm.Deterministic('i', tt.arccos(cosi))
            split = pm.Deterministic('split', xsplit/tt.sin(i))

            # Background treatment
            phi = pm.MvNormal('phi', mu=self.phi_, chol=self.phi_cholesky, testval=self.phi_, shape=len(self.phi_))

            # Construct model
            fit = mod.model([f0, f1, f2, g0, g1, g2, h0, h1, h2, split, i, phi])

            like = pm.Gamma('like', alpha=1., beta=1./fit, observed=self.p)
Esempio n. 13
0
def _get_compiled_theano_functions(N_QUAD_PTS):
    # resonance j and k
    j, k = T.lscalars('jk')
    s = (j - k) / k
    # Planet masses: m1,m2
    m1, m2 = T.dscalars(2)
    mstar = 1
    mu1 = m1 * mstar / (mstar + m1)
    mu2 = m2 * mstar / (mstar + m2)
    eta1 = mstar + m1
    eta2 = mstar + m2
    beta1 = mu1 * T.sqrt(eta1 / mstar) / (mu1 + mu2)
    beta2 = mu2 * T.sqrt(eta2 / mstar) / (mu1 + mu2)

    # Angle variable for averaging over
    psi = T.dvector('psi')

    # Dynamical variables:
    dyvars = T.vector()
    y1, y2, phi1, x1, x2, J1, amd = [dyvars[i] for i in range(7)]

    I1 = (y1 * y1 + x1 * x1) / 2
    I2 = (y2 * y2 + x2 * x2) / 2
    sigma1 = T.arctan2(y1, x1)
    sigma2 = T.arctan2(y2, x2)

    # Quadrature weights
    quad_weights = T.dvector('w')

    # Set l=0
    l = T.constant(0.)
    l1 = l - k * psi / 2
    l2 = l + k * psi / 2
    pomega1 = (1 + s) * l2 - s * l1 - sigma1
    pomega2 = (1 + s) * l2 - s * l1 - sigma2

    Gamma1 = I1
    Gamma2 = I2

    # Resonant semi-major axis ratio
    a20 = T.constant(1.0)
    a10 = (eta1 / eta2)**(1 / 3) * ((j - k) / j)**(2 / 3) * a20
    Lambda20 = beta2 * T.sqrt(a20)
    Lambda10 = beta1 * T.sqrt(a10)
    Ltot = Lambda10 + Lambda20

    Psi0 = 0.5 * k * (Lambda20 - Lambda10)
    Psi = Psi0 - k * (s + 1 / 2) * amd

    L1 = Ltot / 2 + J1 / 2 - Psi / k - s * (I1 + I2)
    L2 = Ltot / 2 + J1 / 2 + Psi / k + (1 + s) * (I1 + I2)

    # Choose z axis along direction of total angular momentum
    r2_by_r1 = (L2 - L1 - Gamma2 + Gamma1) / (L1 + L2 - Gamma1 - Gamma2 - J1)
    rho1 = 0.5 * J1 * (1 + r2_by_r1)
    rho2 = 0.5 * J1 * (1 - r2_by_r1)

    G1 = L1 - Gamma1
    G2 = L2 - Gamma1

    a1 = (L1 / beta1)**2
    e1 = T.sqrt(1 - (1 - (Gamma1 / L1))**2)

    a2 = (L2 / beta2)**2
    e2 = T.sqrt(1 - (1 - (Gamma2 / L2))**2)

    cos_inc1 = 1 - rho1 / G1
    cos_inc2 = 1 - rho2 / G2
    inc1 = T.arccos(cos_inc1)
    inc2 = T.arccos(cos_inc2)
    Omega1 = l - np.pi / 2 - phi1
    Omega2 = l + np.pi / 2 - phi1
    omega1 = pomega1 - Omega1
    omega2 = pomega2 - Omega2

    Hkep = -0.5 * T.sqrt(eta1) * beta1 / a1 - 0.5 * T.sqrt(eta2) * beta2 / a2

    ko = KeplerOp()
    M1 = l1 - pomega1
    M2 = l2 - pomega2
    sinf1, cosf1 = ko(M1, e1 + T.zeros_like(M1))
    sinf2, cosf2 = ko(M2, e2 + T.zeros_like(M2))
    #
    n1 = T.sqrt(eta1 / mstar) * a1**(-3 / 2)
    n2 = T.sqrt(eta2 / mstar) * a2**(-3 / 2)
    Hint_dir, Hint_ind = calc_Hint_components_sinf_cosf(
        a1, a2, e1, e2, inc1, inc2, omega1, omega2, Omega1, Omega2, n1, n2,
        sinf1, cosf1, sinf2, cosf2)
    eps = m1 * m2 / (mu1 + mu2) / T.sqrt(mstar * a20)
    Hpert = (Hint_dir + Hint_ind / mstar).dot(quad_weights)
    Htot = Hkep + eps * Hpert

    #####################################################
    # Set parameters for compiling functions with Theano
    #####################################################

    # Get numerical quadrature nodes and weights
    nodes, weights = np.polynomial.legendre.leggauss(N_QUAD_PTS)

    # Rescale for integration interval from [-1,1] to [-pi,pi]
    nodes = nodes * np.pi
    weights = weights * 0.5

    # 'givens' will fix some parameters of Theano functions compiled below
    givens = [(psi, nodes), (quad_weights, weights)]

    # 'ins' will set the inputs of Theano functions compiled below
    #   Note: 'extra_ins' will be passed as values of object attributes
    #   of the 'ResonanceEquations' class 'defined below
    extra_ins = [m1, m2, j, k]
    ins = [dyvars] + extra_ins

    orbels = [
        a1, e1, inc1, sigma1 * k, Omega1, a2, e2, inc2, sigma2 * k, Omega2
    ]

    #  Conservative flow
    gradHtot = T.grad(Htot, wrt=dyvars)
    hessHtot = theano.gradient.hessian(Htot, wrt=dyvars)
    Jtens = T.as_tensor(np.pad(_get_Omega_matrix(3), (0, 1), 'constant'))
    H_flow_vec = Jtens.dot(gradHtot)
    H_flow_jac = Jtens.dot(hessHtot)

    ##########################
    # Compile Theano functions
    ##########################
    actions_fn = theano.function(
        inputs=ins,
        outputs={'L1':L1,'L1':L2,'Gamma1':Gamma1,'Gamma2':Gamma2,\
                'Q1':rho1,'Q2':rho2,'Psi':Psi,'Psi0':Psi0},
        givens=givens,
        on_unused_input='ignore'
    )
    orbels_fn = theano.function(inputs=ins,
                                outputs=orbels,
                                givens=givens,
                                on_unused_input='ignore')

    Htot_fn = theano.function(inputs=ins,
                              outputs=Htot,
                              givens=givens,
                              on_unused_input='ignore')

    Hpert_fn = theano.function(inputs=ins,
                               outputs=Hpert,
                               givens=givens,
                               on_unused_input='ignore')

    H_flow_vec_fn = theano.function(inputs=ins,
                                    outputs=H_flow_vec,
                                    givens=givens,
                                    on_unused_input='ignore')

    H_flow_jac_fn = theano.function(inputs=ins,
                                    outputs=H_flow_jac,
                                    givens=givens,
                                    on_unused_input='ignore')
    return dict({
        'orbital_elements': orbels_fn,
        'actions': actions_fn,
        'Hamiltonian': Htot_fn,
        'Hpert': Hpert_fn,
        'Hamiltonian_flow': H_flow_vec_fn,
        'Hamiltonian_flow_jacobian': H_flow_jac_fn
    })
Esempio n. 14
0
    def get_camera_image(self, state, camera_name):

        # get camera image
        # do ray-sphere and ray-plane intersections
        # find cube by using 6 planes, throwing away the irrelevant intersections

        # step 1: generate list of rays (1 per pixel)
        # focal_point (3,)
        # ray_dir (px_hor, px_ver, 3)
        # ray_offset (px_hor, px_ver, 3)

        camera = self.cameras[camera_name]
        positions, velocities, rotations = state.positions, state.velocities, state.rotations

        ray_dir = camera["ray_direction"]
        ray_offset = camera["ray_offset"]
        parent = camera["parent"]

        px_ver = ray_dir.shape[0]
        px_hor = ray_dir.shape[1]

        # WORKPOINT

        if parent:
            pid = self.objects[parent]
            # rotate and move the camera according to its parent
            ray_dir = theano_convert_model_to_world_coordinate_no_bias(
                ray_dir, rotations[:, pid, :, :])
            ray_offset = theano_convert_model_to_world_coordinate(
                ray_offset, rotations[:, pid, :, :], positions[:, pid, :])
        else:
            ray_dir = ray_dir[None, :, :, :]
            ray_offset = ray_offset[None, :, :, :]

        # step 2a: intersect the rays with all the spheres
        has_spheres = (0 != len(self.sphere_parent))
        #s_relevant = np.ones(shape=(self.batch_size, px_ver, px_hor, self.sphere_parent.shape[0]))
        if has_spheres:
            s_pos_vectors = positions[:, None, None, self.sphere_parent, :]
            s_rot_matrices = rotations[:, self.sphere_parent, :, :]

            L = s_pos_vectors - ray_offset[:, :, :, None, :]
            tca = T.sum(L * ray_dir[:, :, :, None, :],
                        axis=4)  # L.dotProduct(ray_dir)
            #// if (tca < 0) return false;
            d2 = T.sum(L * L, axis=4) - tca * tca
            r2 = self.sphere_radius**2
            #if (d2 > radius2) return false;
            s_relevant = (tca > 0) * (d2[:, :, :, :] < r2[None, None, None, :])
            float_s_relevant = T.cast(s_relevant, 'float32')
            thc = T.sqrt(
                (r2[None, None, None, :] - float_s_relevant * d2[:, :, :, :]))
            s_t0 = tca - thc
            Phit = ray_offset[:, :, :,
                              None, :] + s_t0[:, :, :, :,
                                              None] * ray_dir[:, :, :, None, :]
            N = (Phit - s_pos_vectors) / self.sphere_radius[None, None,
                                                            None, :, None]
            N = theano_convert_world_to_model_coordinate_no_bias(
                N, s_rot_matrices[:, None, None, :, :, :])

            # tex_y en tex_x in [-1,1]
            s_tex_x = T.arctan2(N[:, :, :, :, 2], N[:, :, :, :, 0]) / np.pi
            s_tex_y = -1. + (2. - eps) * T.arccos(
                T.clip(N[:, :, :, :, 1], -1.0, 1.0)) / np.pi

        # step 2b: intersect the rays with the cubes (cubes=planes)
        # step 2c: intersect the rays with the planes
        has_faces = (0 != len(self.face_parent))
        if has_faces:
            hasparent = [
                i for i, par in enumerate(self.face_parent) if par is not None
            ]
            hasnoparent = [
                i for i, par in enumerate(self.face_parent) if par is None
            ]
            parents = [
                parent for parent in self.face_parent if parent is not None
            ]

            static_fn = numpy_repeat_new_axis(self.face_normal[hasnoparent, :],
                                              self.batch_size)
            static_fp = numpy_repeat_new_axis(self.face_point[hasnoparent, :],
                                              self.batch_size)
            static_ftx = numpy_repeat_new_axis(
                self.face_texture_x[hasnoparent, :], self.batch_size)
            static_fty = numpy_repeat_new_axis(
                self.face_texture_y[hasnoparent, :], self.batch_size)

            if hasparent:
                fn = theano_convert_model_to_world_coordinate_no_bias(
                    self.face_normal[None, hasparent, :],
                    rotations[:, parents, :, :])
                fn = T.concatenate([static_fn, fn], axis=1)
                fp = theano_convert_model_to_world_coordinate(
                    self.face_point[None, hasparent, :],
                    rotations[:, parents, :, :], positions[:, parents, :])
                fp = T.concatenate([static_fp, fp], axis=1)
                ftx = theano_convert_model_to_world_coordinate_no_bias(
                    self.face_texture_x[None, hasparent, :],
                    rotations[:, parents, :, :])
                ftx = T.concatenate([static_ftx, ftx], axis=1)
                fty = theano_convert_model_to_world_coordinate_no_bias(
                    self.face_texture_y[None, hasparent, :],
                    rotations[:, parents, :, :])
                fty = T.concatenate([static_fty, fty], axis=1)
            else:
                fn = static_fn
                fp = static_fp
                ftx = static_ftx
                fty = static_fty

            # reshuffle the face_texture_indexes to match the reshuffling we did above
            face_indices = hasnoparent + hasparent
            face_texture_index = self.face_texture_index[face_indices]
            face_texture_limited = self.face_texture_limited[face_indices]
            face_colors = self.face_colors[face_indices, :]

            denom = T.sum(fn[:, None, None, :, :] * ray_dir[:, :, :, None, :],
                          axis=4)
            p0l0 = fp[:, None, None, :, :] - ray_offset[:, :, :, None, :]
            p_t0 = T.sum(p0l0 * fn[:, None, None, :, :],
                         axis=4) / (denom + 1e-9)

            Phit = ray_offset[:, :, :,
                              None, :] + p_t0[:, :, :, :,
                                              None] * ray_dir[:, :, :, None, :]

            pd = Phit - fp[:, None, None, :, :]
            p_tex_x = T.sum(ftx[:, None, None, :, :] * pd, axis=4)
            p_tex_y = T.sum(fty[:, None, None, :, :] * pd, axis=4)

            # the following only on limited textures
            p_relevant = (p_t0 > 0) * (1 -
                                       (1 - (-1 < p_tex_x) * (p_tex_x < 1) *
                                        (-1 < p_tex_y) *
                                        (p_tex_y < 1)) * face_texture_limited)

            p_tex_x = ((p_tex_x + 1) % 2.) - 1
            p_tex_y = ((p_tex_y + 1) % 2.) - 1

        # step 3: find the closest point of intersection for all objects (z-culling)

        if has_spheres and has_faces:
            relevant = T.concatenate([s_relevant, p_relevant],
                                     axis=3).astype('float32')
            tex_x = T.concatenate([s_tex_x, p_tex_x], axis=3)
            tex_y = T.concatenate([s_tex_y, p_tex_y], axis=3)
            tex_t = np.concatenate(
                [self.sphere_texture_index, face_texture_index], axis=0)
            t = T.concatenate([s_t0, p_t0], axis=3)
        elif has_spheres:
            relevant = s_relevant.astype('float32')
            tex_x = s_tex_x
            tex_y = s_tex_y
            tex_t = self.sphere_texture_index
            t = s_t0
        elif has_faces:
            relevant = p_relevant.astype('float32')
            tex_x = p_tex_x
            tex_y = p_tex_y
            tex_t = face_texture_index
            t = p_t0
        else:
            raise NotImplementedError()

        mint = T.min(t * relevant + (1. - relevant) * np.float32(1e9), axis=3)
        relevant *= (t <= mint[:, :, :, None])  #only use the closest object

        # step 4: go into the object's texture and get the corresponding value (see image transform)
        x_size, y_size = self.textures.shape[1] - 1, self.textures.shape[2] - 1

        tex_x = (tex_x + 1) * x_size / 2.
        tex_y = (tex_y + 1) * y_size / 2.
        x_idx = T.floor(tex_x)
        x_wgh = tex_x - x_idx
        y_idx = T.floor(tex_y)
        y_wgh = tex_y - y_idx

        # if the following are -2,147,483,648 or -9,223,372,036,854,775,808, you have NaN's
        x_idx, y_idx = T.cast(x_idx, 'int64'), T.cast(y_idx, 'int64')

        textures = T.TensorConstant(type=T.ftensor4,
                                    data=self.textures.astype('float32'),
                                    name='textures')

        sample= (   x_wgh  *    y_wgh )[:,:,:,:,None] * textures[tex_t[None,None,None,:],x_idx+1,y_idx+1,:] + \
                (   x_wgh  * (1-y_wgh))[:,:,:,:,None] * textures[tex_t[None,None,None,:],x_idx+1,y_idx  ,:] + \
                ((1-x_wgh) *    y_wgh )[:,:,:,:,None] * textures[tex_t[None,None,None,:],x_idx  ,y_idx+1,:] + \
                ((1-x_wgh) * (1-y_wgh))[:,:,:,:,None] * textures[tex_t[None,None,None,:],x_idx  ,y_idx  ,:]

        # multiply with color of object
        colors = np.concatenate([self.sphere_colors, face_colors], axis=0)
        if np.min(colors) != 1.:  # if the colors are actually used
            sample = colors[None, None, None, :, :] * sample

        # step 5: return this value

        image = T.sum(sample * relevant[:, :, :, :, None], axis=3)

        background_color = camera["background_color"]
        if background_color is not None:
            # find the rays for which no object was relevant. Make them background color
            background = background_color[None, None, None, :] * (
                1 - T.max(relevant[:, :, :, :], axis=3))[:, :, :, None]
            image += background

        # do a dimshuffle to closer match the deep learning conventions
        image = T.unbroadcast(image, 0, 1, 2, 3).dimshuffle(0, 3, 2, 1)
        return image
Esempio n. 15
0
def _get_compiled_theano_functions(N_QUAD_PTS):
    # Planet masses: m1,m2
    m1, m2 = T.dscalars(2)
    mstar = 1
    mu1 = m1 * mstar / (mstar + m1)
    mu2 = m2 * mstar / (mstar + m2)
    eta1 = mstar + m1
    eta2 = mstar + m2
    beta1 = mu1 * T.sqrt(eta1 / mstar) / (mu1 + mu2)
    beta2 = mu2 * T.sqrt(eta2 / mstar) / (mu1 + mu2)
    j, k = T.lscalars('jk')
    s = (j - k) / k

    # Angle variable for averaging over
    psi = T.dvector('psi')

    # Quadrature weights
    quad_weights = T.dvector('w')

    # Dynamical variables:
    Ndof = 3
    Nconst = 1
    dyvars = T.vector()
    s1, s2, phi, I1, I2, Phi, dRtilde = [
        dyvars[i] for i in range(2 * Ndof + Nconst)
    ]

    a20 = T.constant(1.)
    a10 = ((j - k) / j)**(2 / 3) * (eta1 / eta2)**(1 / 3) * a20
    L10 = beta1 * T.sqrt(a10)
    L20 = beta2 * T.sqrt(a20)
    Psi = s * L20 + (1 + s) * L10
    Rtilde = dRtilde - L10 - L20
    ####
    # angles
    ####
    rtilde = T.constant(0.)
    Omega = -1 * rtilde
    l1 = phi + k * (1 + s) * psi + Omega
    l2 = phi + k * s * psi + Omega
    gamma1 = s1 - phi - Omega
    gamma2 = s2 - phi - Omega
    q1 = 0.5 * np.pi - Omega
    q2 = -0.5 * np.pi - Omega

    pomega1 = -1 * gamma1
    pomega2 = -1 * gamma2
    Omega1 = -1 * q1
    Omega2 = -1 * q2
    omega1 = pomega1 - Omega1
    omega2 = pomega2 - Omega2

    ###
    # actions
    ###
    Gamma1 = I1
    Gamma2 = I2
    L1 = Psi / k - s * (I1 + I2) - s * Phi
    L2 = -1 * Psi / k + (1 + s) * (I1 + I2) + (1 + s) * Phi
    Cz = -1 * Rtilde

    R = L1 + L2 - Gamma1 - Gamma2 - Cz
    G1 = L1 - Gamma1
    G2 = L2 - Gamma2

    r2_by_r1 = (L2 - L1 - Gamma2 + Gamma1) / (L1 + L2 - Gamma1 - Gamma2 - R)
    rho1 = 0.5 * R * (1 + r2_by_r1)
    rho2 = 0.5 * R * (1 - r2_by_r1)

    a1 = (L1 / beta1)**2
    e1 = T.sqrt(1 - (1 - (Gamma1 / L1))**2)

    a2 = (L2 / beta2)**2
    e2 = T.sqrt(1 - (1 - (Gamma2 / L2))**2)

    cos_inc1 = 1 - rho1 / G1
    cos_inc2 = 1 - rho2 / G2
    inc1 = T.arccos(cos_inc1)
    inc2 = T.arccos(cos_inc2)

    Hkep = -0.5 * T.sqrt(eta1) * beta1 / a1 - 0.5 * T.sqrt(eta2) * beta2 / a2

    ko = KeplerOp()
    M1 = l1 - pomega1
    M2 = l2 - pomega2
    sinf1, cosf1 = ko(M1, e1 + T.zeros_like(M1))
    sinf2, cosf2 = ko(M2, e2 + T.zeros_like(M2))
    #
    n1 = T.sqrt(eta1 / mstar) * a1**(-3 / 2)
    n2 = T.sqrt(eta2 / mstar) * a2**(-3 / 2)
    Hint_dir, Hint_ind, r1, r2, v1, v2 = calc_Hint_components_sinf_cosf(
        a1, a2, e1, e2, inc1, inc2, omega1, omega2, Omega1, Omega2, n1, n2,
        sinf1, cosf1, sinf2, cosf2)
    eps = m1 * m2 / (mu1 + mu2) / T.sqrt(mstar)
    Hpert = (Hint_dir + Hint_ind / mstar)
    Hpert_av = Hpert.dot(quad_weights)
    Htot = Hkep + eps * Hpert_av

    #####################################################
    # Set parameters for compiling functions with Theano
    #####################################################

    # Get numerical quadrature nodes and weights
    nodes, weights = np.polynomial.legendre.leggauss(N_QUAD_PTS)

    # Rescale for integration interval from [-1,1] to [-pi,pi]
    nodes = nodes * np.pi
    weights = weights * 0.5

    # 'givens' will fix some parameters of Theano functions compiled below
    givens = [(psi, nodes), (quad_weights, weights)]

    # 'ins' will set the inputs of Theano functions compiled below
    #   Note: 'extra_ins' will be passed as values of object attributes
    #   of the 'ResonanceEquations' class 'defined below
    extra_ins = [m1, m2, j, k]
    ins = [dyvars] + extra_ins
    orbels = [a1, e1, inc1, k * s1, a2, e2, inc2, k * s2, phi, Omega]
    orbels_dict = dict(
        zip([
            'a1', 'e1', 'inc1', 'theta1', 'a2', 'e2', 'inc2', 'theta2', 'phi'
        ], orbels))
    actions = [L1, L2, Gamma1, Gamma2, rho1, rho2]
    actions_dict = dict(
        zip(['L1', 'L2', 'Gamma1', 'Gamma2', 'Q1', 'Q2'], actions))

    #  Conservative flow
    gradHtot = T.grad(Htot, wrt=dyvars)
    hessHtot = theano.gradient.hessian(Htot, wrt=dyvars)
    Jtens = T.as_tensor(
        np.pad(_get_Omega_matrix(Ndof), (0, Nconst), 'constant'))
    H_flow_vec = Jtens.dot(gradHtot)
    H_flow_jac = Jtens.dot(hessHtot)

    ##########################
    # Compile Theano functions
    ##########################
    orbels_fn = theano.function(inputs=ins,
                                outputs=orbels_dict,
                                givens=givens,
                                on_unused_input='ignore')
    actions_fn = theano.function(inputs=ins,
                                 outputs=actions_dict,
                                 givens=givens,
                                 on_unused_input='ignore')
    Rtilde_fn = theano.function(inputs=ins,
                                outputs=Rtilde,
                                givens=givens,
                                on_unused_input='ignore')

    Htot_fn = theano.function(inputs=ins,
                              outputs=Htot,
                              givens=givens,
                              on_unused_input='ignore')

    Hpert_fn = theano.function(inputs=ins,
                               outputs=Hpert_av,
                               givens=givens,
                               on_unused_input='ignore')

    Hpert_components_fn = theano.function(
        inputs=ins,
        outputs=[Hint_dir.dot(quad_weights),
                 Hint_ind.dot(quad_weights)],
        givens=givens,
        on_unused_input='ignore')

    H_flow_vec_fn = theano.function(inputs=ins,
                                    outputs=H_flow_vec,
                                    givens=givens,
                                    on_unused_input='ignore')

    H_flow_jac_fn = theano.function(inputs=ins,
                                    outputs=H_flow_jac,
                                    givens=givens,
                                    on_unused_input='ignore')

    return dict({
        'orbital_elements': orbels_fn,
        'actions': actions_fn,
        'Rtilde': Rtilde_fn,
        'Hamiltonian': Htot_fn,
        'Hpert': Hpert_fn,
        'Hpert_components': Hpert_components_fn,
        'Hamiltonian_flow': H_flow_vec_fn,
        'Hamiltonian_flow_jacobian': H_flow_jac_fn
    })
Esempio n. 16
0
  return vals[0], variance # T.maximum(T.var(s, axis=0), eps) # Should eps be here?!
  # return T.mean(s, axis=0), variance # T.maximum(T.var(s, axis=0), eps) # Should eps be here?!

def bound_loss(x, tnp = np) :
  eps = 1e-9
  loss = tnp.maximum(tnp.maximum(eps,x-1), tnp.maximum(eps,-x)) + eps
  return tnp.maximum(loss, eps) + eps

params_plus1 = shared(np.random.rand(nbatch, 2))
params_plus2 = shared(np.random.rand(nbatch, 2))

a = Print('invplus(x, params_plus1')(invplus(x, params_plus1))
b = Print('invplus(x, params_plus2')(invplus(y, params_plus2))

eps = 1e-9
a2 = T.arccos(T.clip(a, eps, 1))
b2 = T.arcsin(T.clip(b, eps, 1))
bl1 = bound_loss(a, tnp = T)
bl2 = bound_loss(b, tnp = T)

theta1_group = []
theta2_group = []
theta3_group = []
theta1_group.append(a2[:, 0])
theta1_group.append(b2[:, 0])

delta_group = []
delta_group.append(a2[:, 1])
delta_group.append(b2[:, 1])
delta_single, delta_var = singular(delta_group, name = 'delta')
Esempio n. 17
0
    def __graph_output(self, epoch):

        self.update_learning_stats_fn()

        ########  Learning statistic associated with optimized parameters
        for param in self.params:

            last_update = self.last_batch_update[param]
            this_update = self.this_batch_update[param]
            init = self.init[param]
            sp = self.subplots[param]
            name = str(param)

            if param.ndim == 1:
                data = param.get_value()
                p005, median, p995 = np.percentile(data, [0.5, 50, 99.5])
                sp[0, 0].hist(remove=True, x=data, bins=35)
                sp[0, 0].set_title('{} at epoch {}'.format(param, epoch),
                                   fontsize=10)
                sp[1, 0].add_point(p005=(epoch, p005),
                                   median=(epoch, median),
                                   p995=(epoch, p995),
                                   std=(epoch, data.std()))

                data = (param - init).eval()
                p005, median, p995 = np.percentile(data, [0.5, 50, 99.5])
                sp[0, 1].hist(remove=True, x=data, bins=35)
                sp[0, 1].set_title('{}-initial'.format(param), fontsize=10)
                sp[1, 1].add_point(p005=(epoch, p005),
                                   median=(epoch, median),
                                   p995=(epoch, p995),
                                   std=(epoch, data.std()))

                data = this_update.get_value()
                p005, median, p995 = np.percentile(data, [0.5, 50, 99.5])
                sp[0, 2].hist(remove=True, x=data, bins=35)
                sp[0, 2].set_title('{} gradient update'.format(param),
                                   fontsize=10)
                sp[1, 2].add_point(p005=(epoch, p005),
                                   median=(epoch, median),
                                   p995=(epoch, p995),
                                   std=(epoch, data.std()))

            elif param.ndim > 1:

                if param.ndim == 2:
                    param = param.T
                    last_update = last_update.T
                    this_update = this_update.T
                    init = init.T

                param = param.flatten(2)
                last_update = last_update.flatten(2)
                this_update = this_update.flatten(2)
                init = init.flatten(2)

                # Norms
                nrm = T.sqrt((param**2).sum(1))
                data = nrm.eval()
                p005, median, p995 = np.percentile(data, [0.5, 50, 99.5])
                sp[0, 0].hist(remove=True, x=data, bins=35)
                sp[0, 0].set_title(
                    r'$\Vert w_i \Vert \/ i \in [1,{}]$ at epoch {}'.format(
                        len(data), epoch),
                    fontsize=10)
                sp[1, 0].add_point(p005=(epoch, p005),
                                   median=(epoch, median),
                                   p995=(epoch, p995),
                                   std=(epoch, data.std()))

                # Orthonormality
                param_nrm = param / nrm[:, None]
                data = T.dot(param_nrm, param_nrm.T).flatten().eval()
                p005, median, p995 = np.percentile(data, [0.5, 50, 99.5])
                sp[0, 1].hist(remove=True, x=data, bins=60)
                sp[0, 1].set_yscale('log', nonposy='clip')
                sp[0, 1].set_title(
                    r'$ {{ \frac{{ {{w_i}}^\intercal w_j }}{{ \Vert w_i \Vert \Vert w_j \Vert }} }} {{\vert}}_{{(t={})}}  \/ i,j \in [1,{}]  $'
                    .format(epoch, int(sqrt(len(data)))),
                    fontsize=10)
                sp[1, 1].add_point(p005=(epoch, p005),
                                   median=(epoch, median),
                                   p995=(epoch, p995),
                                   std=(epoch, data.std()))

                # Rotations with respect to initial state
                cos = (param * init).sum(1)
                nrm = T.sqrt((param**2).sum(1))
                nrm_init = T.sqrt((init**2).sum(1))
                fac = cast_x(180. / np.pi)
                data = (T.arccos(cos / (nrm * nrm_init)) *
                        fac).flatten().eval()
                p005, median, p995 = np.percentile(data, [0.5, 50, 99.5])
                sp[0, 2].hist(remove=True, x=data, bins=35)
                sp[0, 2].set_title(
                    r'$ \measuredangle ( w^{{(t={})}}_i, w^{{(t=0)}}_i ) \/ i \in [1,{}] $'
                    .format(epoch, len(data)),
                    fontsize=10)
                sp[1, 2].add_point(p005=(epoch, p005),
                                   median=(epoch, median),
                                   p995=(epoch, p995),
                                   std=(epoch, data.std()))

                # Update norm
                data = T.sqrt((this_update**2).sum(1)).eval()
                p005, median, p995 = np.percentile(data, [0.5, 50, 99.5])
                sp[0, 3].hist(remove=True, x=data, bins=35)
                sp[0, 3].set_title(
                    r'$\Vert u_i \Vert \/ i \in [1,{}]$ at epoch {}'.format(
                        len(data), epoch),
                    fontsize=10)
                sp[1, 3].add_point(p005=(epoch, p005),
                                   median=(epoch, median),
                                   p995=(epoch, p995),
                                   std=(epoch, data.std()))

                # Update rotation with respect to weight vectors
                cos = (param * this_update).sum(1)
                nrm = T.sqrt((param**2).sum(1))
                nrm_init = T.sqrt((this_update**2).sum(1))
                fac = cast_x(180. / np.pi)
                data = (T.arccos(cos / (nrm * nrm_init)) *
                        fac).flatten().eval()
                p005, median, p995 = np.percentile(data, [0.5, 50, 99.5])
                try:
                    sp[0, 4].hist(remove=True, x=data, bins=35)
                except:
                    print(param)
                    print(data.shape)
                    raise
                sp[0, 4].set_title(
                    r'$ \measuredangle ( w^{{(t={})}}_i, u^{{(t={})}}_i ) \/ i \in [1,{}] $'
                    .format(epoch, epoch, len(data)),
                    fontsize=10)
                sp[1, 4].add_point(p005=(epoch, p005),
                                   median=(epoch, median),
                                   p995=(epoch, p995),
                                   std=(epoch, data.std()))

                # Update rotation if this update with respect to the last
                cos = (this_update * last_update).sum(1)
                nrm_this = T.sqrt((last_update**2).sum(1))
                nrm_last = T.sqrt((this_update**2).sum(1))
                fac = cast_x(180. / np.pi)
                data = (T.arccos(cos / (nrm_this * nrm_last)) *
                        fac).flatten().eval()
                p005, median, p995 = np.percentile(data, [0.5, 50, 99.5])
                sp[0, 5].hist(remove=True, x=data, bins=35)
                sp[0, 5].set_title(
                    r'$ \measuredangle ( u^{{(t={})}}_i, u^{{(t={})}}_i ) \/ i \in [1,{}] $'
                    .format(epoch, epoch - 1, len(data)),
                    fontsize=10)
                sp[1, 5].add_point(p005=(epoch, p005),
                                   median=(epoch, median),
                                   p995=(epoch, p995),
                                   std=(epoch, data.std()))

            else:
                continue

            sp.savefig(join(
                self.output_path, '{}_{}_learning_stats_{}.png'.format({
                    0:
                    'unsupervised',
                    1:
                    'supervised'
                }[self.supervised], self.model_id, name)),
                       dpi=100)

        ########  Learning statistic associated with optimized parameters
        if self.debug_nodes:
            outputs = self.debug_fn()
            for (name, node), data in zip(list(self.debug_nodes.items()),
                                          outputs):
                sp = self.subplots[node]
                data = data.flatten()
                nonzeros = float((data != 0).mean())
                p005, median, p995 = np.percentile(data, [0.5, 50, 99.5])
                sp[0, 0].hist(remove=True, x=data, bins=60)
                sp[0, 0].set_yscale('log', nonposy='clip')
                sp[0, 0].set_title('{} at t={}'.format(name, epoch),
                                   fontsize=6)
                sp[1, 0].add_point(p005=(epoch, p005),
                                   median=(epoch, median),
                                   p995=(epoch, p995),
                                   std=(epoch, data.std()),
                                   nonzero=(epoch, nonzeros))
                sp[1,
                   0].set_title('Non-zero = {:0.4f}%'.format(nonzeros * 100),
                                fontsize=8)
                sp.savefig(join(self.output_path, name + '.png'), dpi=100)
Esempio n. 18
0
    def __init__(self,
                 period=None,
                 a=None,
                 t0=None,
                 t_periastron=None,
                 incl=None,
                 b=None,
                 duration=None,
                 ecc=None,
                 omega=None,
                 Omega=None,
                 m_planet=0.0,
                 m_star=None,
                 r_star=None,
                 rho_star=None,
                 m_planet_units=None,
                 rho_star_units=None,
                 model=None,
                 contact_points_kwargs=None,
                 **kwargs):
        add_citations_to_model(self.__citations__, model=model)

        self.kepler_op = KeplerOp(**kwargs)

        # Parameters
        if m_planet_units is not None:
            warnings.warn(
                "The `m_planet_units` argument has been deprecated. "
                "Use `with_unit` instead.",
                DeprecationWarning,
            )
            m_planet = with_unit(m_planet, m_planet_units)
        if rho_star_units is not None:
            warnings.warn(
                "The `rho_star_units` argument has been deprecated. "
                "Use `with_unit` instead.",
                DeprecationWarning,
            )
            rho_star = with_unit(rho_star, rho_star_units)
        inputs = _get_consistent_inputs(a, period, rho_star, r_star, m_star,
                                        m_planet)
        (
            self.a,
            self.period,
            self.rho_star,
            self.r_star,
            self.m_star,
            self.m_planet,
        ) = inputs
        self.m_total = self.m_star + self.m_planet

        self.n = 2 * np.pi / self.period
        self.a_star = self.a * self.m_planet / self.m_total
        self.a_planet = -self.a * self.m_star / self.m_total

        self.K0 = self.n * self.a / self.m_total

        # Set up the contact points calculation
        if contact_points_kwargs is None:
            contact_points_kwargs = dict()

        if Omega is None:
            self.Omega = None
        else:
            self.Omega = tt.as_tensor_variable(Omega)
            self.cos_Omega = tt.cos(self.Omega)
            self.sin_Omega = tt.sin(self.Omega)

        # Eccentricity
        self.contact_points_op = ContactPointsOp(**contact_points_kwargs)
        if ecc is None:
            self.ecc = None
            self.M0 = 0.5 * np.pi + tt.zeros_like(self.n)
            incl_factor = 1
        else:
            self.ecc = tt.as_tensor_variable(ecc)
            if omega is None:
                raise ValueError("both e and omega must be provided")
            self.omega = tt.as_tensor_variable(omega)

            self.cos_omega = tt.cos(self.omega)
            self.sin_omega = tt.sin(self.omega)

            opsw = 1 + self.sin_omega
            E0 = 2 * tt.arctan2(
                tt.sqrt(1 - self.ecc) * self.cos_omega,
                tt.sqrt(1 + self.ecc) * opsw,
            )
            self.M0 = E0 - self.ecc * tt.sin(E0)

            ome2 = 1 - self.ecc**2
            self.K0 /= tt.sqrt(ome2)
            incl_factor = (1 + self.ecc * self.sin_omega) / ome2

        # The Jacobian for the transform cos(i) -> b
        self.dcosidb = incl_factor * self.r_star / self.a

        if b is not None:
            if incl is not None or duration is not None:
                raise ValueError(
                    "only one of 'incl', 'b', and 'duration' can be given")
            self.b = tt.as_tensor_variable(b)
            self.cos_incl = self.dcosidb * self.b
            self.incl = tt.arccos(self.cos_incl)
        elif incl is not None:
            if duration is not None:
                raise ValueError(
                    "only one of 'incl', 'b', and 'duration' can be given")
            self.incl = tt.as_tensor_variable(incl)
            self.cos_incl = tt.cos(self.incl)
            self.b = self.cos_incl / self.dcosidb
        elif duration is not None:
            if self.ecc is None:
                raise ValueError(
                    "fitting with duration only works for eccentric orbits")
            self.duration = tt.as_tensor_variable(to_unit(duration, u.day))
            c = tt.sin(np.pi * self.duration * incl_factor / self.period)
            c2 = c * c
            aor = self.a_planet / self.r_star
            esinw = self.ecc * self.sin_omega
            self.b = tt.sqrt(
                (aor**2 * c2 - 1) / (c2 * esinw**2 + 2 * c2 * esinw + c2 -
                                     self.ecc**4 + 2 * self.ecc**2 - 1))
            self.b *= 1 - self.ecc**2
            self.cos_incl = self.dcosidb * self.b
            self.incl = tt.arccos(self.cos_incl)
        else:
            zla = tt.zeros_like(self.a)
            self.incl = 0.5 * np.pi + zla
            self.cos_incl = zla
            self.b = zla

        if t0 is not None and t_periastron is not None:
            raise ValueError("you can't define both t0 and t_periastron")
        if t0 is None and t_periastron is None:
            t0 = tt.zeros_like(self.period)

        if t0 is None:
            self.t_periastron = tt.as_tensor_variable(t_periastron)
            self.t0 = self.t_periastron + self.M0 / self.n
        else:
            self.t0 = tt.as_tensor_variable(t0)
            self.t_periastron = self.t0 - self.M0 / self.n

        self.tref = self.t_periastron - self.t0

        self.sin_incl = tt.sin(self.incl)
##### print(f"start position of real C_alpha trace within loaded protein: {rand}")
##### cas_real = cas_real[rand: rand+n] #####################################
# Vertical stack to concatenate list of lists of coordinates into np array
cas_real_array = np.vstack(cas_real)
cas_real_array_centered = center(cas_real_array)
M_obs = cas_real_array_centered
###########

with pm.Model() as inner_model:

    ### 1. prior over mean M
    x1 = pm.Uniform("x1", 0, 1, shape=n)
    theta = 2 * np.pi * x1

    x2 = pm.Uniform("x2", 0, 1, shape=n)
    phi = tt.arccos(1 - 2 * x2)

    x = tt.sin(phi) * tt.cos(theta)
    y = tt.sin(phi) * tt.sin(theta)
    z = tt.cos(phi)

    # stack the 3 coordinate-column-vectors horizontally to get the n*3 matrix of all C_alpha positions
    # adjust length of unit vectors to 3.8 A (average C_alpha distance)
    cas = pm.Deterministic("cas", 3.8 * tt.stack([x, y, z], axis=1))

    # make a C_alpha-trace out of the C_alpha positions lying on the sphere with radius 3.8
    cas_trace = tt.extra_ops.cumsum(cas, axis=0)

    # Move center of mass of C_alpha trace to origin
    M = cas_trace - tt.mean(cas_trace, axis=0)
Esempio n. 20
0
def acos(x):
  return T.arccos(x)
Esempio n. 21
0
    logP_inner = pm.Uniform("logP_inner",
                            lower=0,
                            upper=np.log(50.0),
                            testval=np.log(34.879))  # days
    P_inner = pm.Deterministic("P_inner", tt.exp(logP_inner))

    e_inner = pm.Uniform("e_inner", lower=0, upper=1, testval=0.63)

    omega_inner = Angle("omega_inner", testval=1.42)  # omega_Aa
    Omega_inner = Angle("Omega_inner", testval=1.99)

    cos_incl_inner = pm.Uniform("cos_incl_inner",
                                lower=-1.0,
                                upper=1.0,
                                testval=0.67)
    incl_inner = pm.Deterministic("incl_inner", tt.arccos(cos_incl_inner))

    MAb = pm.Normal("MAb", mu=0.29, sd=0.5)  # solar masses

    t_periastron_inner = pm.Uniform("t_periastron_inner",
                                    lower=4050.0,
                                    upper=4100.0,
                                    testval=4073.0)  # + 2400000 + jd0

    orbit_inner = xo.orbits.KeplerianOrbit(
        a=a_inner * au_to_R_sun,
        period=P_inner,
        ecc=e_inner,
        t_periastron=t_periastron_inner,
        omega=omega_inner,
        Omega=Omega_inner,
def _get_compiled_theano_functions():
    # Planet masses: m1,m2
    m1, m2 = T.dscalars(2)
    mstar = 1
    mu1 = m1 * mstar / (mstar + m1)
    mu2 = m2 * mstar / (mstar + m2)
    eta1 = mstar + m1
    eta2 = mstar + m2
    beta1 = mu1 * T.sqrt(eta1 / mstar) / (mu1 + mu2)
    beta2 = mu2 * T.sqrt(eta2 / mstar) / (mu1 + mu2)

    # Dynamical variables:
    dyvars = T.vector()
    l1, l2, y1, y2, Omega, L1, L2, x1, x2, Rtilde = [
        dyvars[i] for i in range(10)
    ]

    Gamma1 = 0.5 * (x1 * x1 + y1 * y1)
    Gamma2 = 0.5 * (x2 * x2 + y2 * y2)
    gamma1 = T.arctan2(y1, x1)
    gamma2 = T.arctan2(y2, x2)

    Cz = -1 * Rtilde

    R1 = Gamma1
    R2 = Gamma2
    R = L1 + L2 - R1 - R2 - Cz
    G1 = L1 - Gamma1
    G2 = L2 - Gamma2

    r2_by_r1 = (L2 - L1 - R2 + R1) / (L1 + L2 - R1 - R2 - R)
    rho1 = 0.5 * R * (1 + r2_by_r1)
    rho2 = 0.5 * R * (1 - r2_by_r1)

    a1 = (L1 / beta1)**2
    e1 = T.sqrt(1 - (1 - (Gamma1 / L1))**2)

    a2 = (L2 / beta2)**2
    e2 = T.sqrt(1 - (1 - (Gamma2 / L2))**2)

    cos_inc1 = 1 - rho1 / G1
    cos_inc2 = 1 - rho2 / G2
    inc1 = T.arccos(cos_inc1)
    inc2 = T.arccos(cos_inc2)

    l1_r = l1 - Omega
    l2_r = l2 - Omega

    Omega1_r = T.constant(np.pi / 2) - Omega
    Omega2_r = Omega1_r - T.constant(np.pi)

    pomega1 = -1 * gamma1
    pomega2 = -1 * gamma2

    pomega1_r = pomega1 - Omega
    pomega2_r = pomega2 - Omega

    omega1 = pomega1_r - Omega1_r
    omega2 = pomega2_r - Omega2_r

    Hkep = -0.5 * T.sqrt(eta1) * beta1 / a1 - 0.5 * T.sqrt(eta2) * beta2 / a2

    ko = KeplerOp()
    M1 = l1_r - pomega1_r
    M2 = l2_r - pomega2_r
    sinf1, cosf1 = ko(M1, e1 + T.zeros_like(M1))
    sinf2, cosf2 = ko(M2, e2 + T.zeros_like(M2))
    #
    n1 = T.sqrt(eta1 / mstar) * a1**(-3 / 2)
    n2 = T.sqrt(eta2 / mstar) * a2**(-3 / 2)
    Hint_dir, Hint_ind, r1, r2, v1, v2 = calc_Hint_components_sinf_cosf(
        a1, a2, e1, e2, inc1, inc2, omega1, omega2, Omega1_r, Omega2_r, n1, n2,
        sinf1, cosf1, sinf2, cosf2)
    eps = m1 * m2 / (mu1 + mu2) / T.sqrt(mstar)
    Hpert = (Hint_dir + Hint_ind / mstar)
    Htot = Hkep + eps * Hpert

    #####################################################
    # Set parameters for compiling functions with Theano
    #####################################################

    # 'ins' will set the inputs of Theano functions compiled below
    #   Note: 'extra_ins' will be passed as values of object attributes
    #   of the 'ResonanceEquations' class 'defined below
    extra_ins = [m1, m2]
    givens = []
    ins = [dyvars] + extra_ins
    orbels = [
        a1, e1, inc1, l1_r, pomega1_r, Omega1_r, a2, e2, inc2, l2_r, pomega2_r,
        Omega2_r
    ]
    orbels_dict = dict(
        zip([
            'a1', 'e1', 'inc1', 'l1', 'pomega1', 'Omega1', 'a2', 'e2', 'inc2',
            'l2', 'pomega2', 'Omega2'
        ], orbels))
    actions = [L1, L2, Gamma1, Gamma2, rho1, rho2]
    actions_dict = dict(
        zip(['L1', 'L2', 'Gamma1', 'Gamma2', 'Q1', 'Q2'], actions))
    #  Conservative flow
    gradHtot = T.grad(Htot, wrt=dyvars)
    hessHtot = theano.gradient.hessian(Htot, wrt=dyvars)
    Jtens = T.as_tensor(_get_Omega_matrix(5))
    H_flow_vec = Jtens.dot(gradHtot)
    H_flow_jac = Jtens.dot(hessHtot)

    ##########################
    # Compile Theano functions
    ##########################
    orbels_fn = theano.function(inputs=ins,
                                outputs=orbels_dict,
                                givens=givens,
                                on_unused_input='ignore')
    actions_fn = theano.function(inputs=ins,
                                 outputs=actions_dict,
                                 givens=givens,
                                 on_unused_input='ignore')
    rv1_fn = theano.function(inputs=ins,
                             outputs=r1 + v1,
                             givens=givens,
                             on_unused_input='ignore')
    rv2_fn = theano.function(inputs=ins,
                             outputs=r2 + v2,
                             givens=givens,
                             on_unused_input='ignore')

    Htot_fn = theano.function(inputs=ins,
                              outputs=Htot,
                              givens=givens,
                              on_unused_input='ignore')

    Hpert_fn = theano.function(inputs=ins,
                               outputs=Hpert,
                               givens=givens,
                               on_unused_input='ignore')

    Hpert_components_fn = theano.function(inputs=ins,
                                          outputs=[Hint_dir, Hint_ind],
                                          givens=givens,
                                          on_unused_input='ignore')

    H_flow_vec_fn = theano.function(inputs=ins,
                                    outputs=H_flow_vec,
                                    givens=givens,
                                    on_unused_input='ignore')

    H_flow_jac_fn = theano.function(inputs=ins,
                                    outputs=H_flow_jac,
                                    givens=givens,
                                    on_unused_input='ignore')
    return dict({
        'orbital_elements': orbels_fn,
        'actions': actions_fn,
        'Hamiltonian': Htot_fn,
        'Hpert': Hpert_fn,
        'Hpert_components': Hpert_components_fn,
        'Hamiltonian_flow': H_flow_vec_fn,
        'Hamiltonian_flow_jacobian': H_flow_jac_fn,
        'positions_and_velocities1': rv1_fn,
        'positions_and_velocities2': rv2_fn
    })
Esempio n. 23
0
def cos_loss(x1, x2, types, angular=False):
    cossim = cos(x1, x2)
    if angular:
        cossim = 1 - 2 * T.arccos(cossim - 1e-6) / np.pi
    cos_cost = T.switch(types, 1 - cossim, cossim)
    return cos_cost
Esempio n. 24
0
    def __init__(self,
                 period=None,
                 a=None,
                 t0=None,
                 t_periastron=None,
                 incl=None,
                 b=None,
                 duration=None,
                 ecc=None,
                 omega=None,
                 sin_omega=None,
                 cos_omega=None,
                 Omega=None,
                 m_planet=0.0,
                 m_star=None,
                 r_star=None,
                 rho_star=None,
                 ror=None,
                 m_planet_units=None,
                 rho_star_units=None,
                 model=None,
                 contact_points_kwargs=None,
                 **kwargs):
        add_citations_to_model(self.__citations__, model=model)

        self.jacobians = defaultdict(lambda: defaultdict(None))

        daordtau = None
        if ecc is None and duration is not None:
            if r_star is None:
                r_star = as_tensor_variable(1.0)
            if b is None:
                raise ValueError(
                    "'b' must be provided for a circular orbit with a "
                    "'duration'")
            aor, daordtau = get_aor_from_transit_duration(duration,
                                                          period,
                                                          b,
                                                          ror=ror)
            a = r_star * aor
            duration = None

        # Parameters
        if m_planet_units is not None:
            warnings.warn(
                "The `m_planet_units` argument has been deprecated. "
                "Use `with_unit` instead.",
                DeprecationWarning,
            )
            m_planet = with_unit(m_planet, m_planet_units)
        if rho_star_units is not None:
            warnings.warn(
                "The `rho_star_units` argument has been deprecated. "
                "Use `with_unit` instead.",
                DeprecationWarning,
            )
            rho_star = with_unit(rho_star, rho_star_units)
        inputs = _get_consistent_inputs(a, period, rho_star, r_star, m_star,
                                        m_planet)
        (
            self.a,
            self.period,
            self.rho_star,
            self.r_star,
            self.m_star,
            self.m_planet,
        ) = inputs
        self.m_total = self.m_star + self.m_planet

        self.n = 2 * np.pi / self.period
        self.a_star = self.a * self.m_planet / self.m_total
        self.a_planet = -self.a * self.m_star / self.m_total

        # Track the Jacobian between the duration and a
        if daordtau is not None:
            dadtau = self.r_star * daordtau
            self.jacobians["duration"]["a"] = dadtau
            self.jacobians["duration"]["a_star"] = (dadtau * self.m_planet /
                                                    self.m_total)
            self.jacobians["duration"]["a_planet"] = (-dadtau * self.m_star /
                                                      self.m_total)

            # rho = 3 * pi * (a/R)**3 / (G * P**2)
            # -> drho / d(a/R) = 9 * pi * (a/R)**2 / (G * P**2)
            self.jacobians["duration"]["rho_star"] = (
                9 * np.pi * (self.a / self.r_star)**2 * daordtau *
                gcc_per_sun / (G_grav * self.period**2))

        self.K0 = self.n * self.a / self.m_total

        # Set up the contact points calculation
        if contact_points_kwargs is None:
            self.contact_points = contact_points
        else:
            self.contact_points = ContactPoints(**contact_points_kwargs)

        if Omega is None:
            self.Omega = None
        else:
            self.Omega = as_tensor_variable(Omega)
            self.cos_Omega = tt.cos(self.Omega)
            self.sin_Omega = tt.sin(self.Omega)

        # Eccentricity
        if ecc is None:
            self.ecc = None
            self.M0 = 0.5 * np.pi + tt.zeros_like(self.n)
            incl_factor = 1
        else:
            self.ecc = as_tensor_variable(ecc)
            if omega is not None:
                if sin_omega is not None and cos_omega is not None:
                    raise ValueError(
                        "either 'omega' or 'sin_omega' and 'cos_omega' can be "
                        "provided")
                self.omega = as_tensor_variable(omega)
                self.cos_omega = tt.cos(self.omega)
                self.sin_omega = tt.sin(self.omega)
            elif sin_omega is not None and cos_omega is not None:
                self.cos_omega = as_tensor_variable(cos_omega)
                self.sin_omega = as_tensor_variable(sin_omega)
                self.omega = tt.arctan2(self.sin_omega, self.cos_omega)

            else:
                raise ValueError("both e and omega must be provided")

            opsw = 1 + self.sin_omega
            E0 = 2 * tt.arctan2(
                tt.sqrt(1 - self.ecc) * self.cos_omega,
                tt.sqrt(1 + self.ecc) * opsw,
            )
            self.M0 = E0 - self.ecc * tt.sin(E0)

            ome2 = 1 - self.ecc**2
            self.K0 /= tt.sqrt(ome2)
            incl_factor = (1 + self.ecc * self.sin_omega) / ome2

        # The Jacobian for the transform cos(i) -> b
        self.dcosidb = self.jacobians["b"]["cos_incl"] = (incl_factor *
                                                          self.r_star / self.a)

        if b is not None:
            if incl is not None or duration is not None:
                raise ValueError(
                    "only one of 'incl', 'b', and 'duration' can be given")
            self.b = as_tensor_variable(b)
            self.cos_incl = self.dcosidb * self.b
            self.incl = tt.arccos(self.cos_incl)
        elif incl is not None:
            if duration is not None:
                raise ValueError(
                    "only one of 'incl', 'b', and 'duration' can be given")
            self.incl = as_tensor_variable(incl)
            self.cos_incl = tt.cos(self.incl)
            self.b = self.cos_incl / self.dcosidb
        elif duration is not None:
            if self.ecc is None:
                raise ValueError(
                    "fitting with duration only works for eccentric orbits")
            self.duration = as_tensor_variable(to_unit(duration, u.day))
            c = tt.sin(np.pi * self.duration * incl_factor / self.period)
            c2 = c * c
            aor = self.a_planet / self.r_star
            esinw = self.ecc * self.sin_omega
            self.b = tt.sqrt(
                (aor**2 * c2 - 1) / (c2 * esinw**2 + 2 * c2 * esinw + c2 -
                                     self.ecc**4 + 2 * self.ecc**2 - 1))
            self.b *= 1 - self.ecc**2
            self.cos_incl = self.dcosidb * self.b
            self.incl = tt.arccos(self.cos_incl)
        else:
            zla = tt.zeros_like(self.a)
            self.incl = 0.5 * np.pi + zla
            self.cos_incl = zla
            self.b = zla

        if t0 is not None and t_periastron is not None:
            raise ValueError("you can't define both t0 and t_periastron")
        if t0 is None and t_periastron is None:
            t0 = tt.zeros_like(self.period)

        if t0 is None:
            self.t_periastron = as_tensor_variable(t_periastron)
            self.t0 = self.t_periastron + self.M0 / self.n
        else:
            self.t0 = as_tensor_variable(t0)
            self.t_periastron = self.t0 - self.M0 / self.n

        self.tref = self.t_periastron - self.t0

        self.sin_incl = tt.sin(self.incl)
Esempio n. 25
0
    def __init__(self,
                 period=None, a=None, t0=0.0,
                 incl=None, b=None, duration=None,
                 ecc=None, omega=None, m_planet=0.0,
                 m_star=None, r_star=None, rho_star=None,
                 m_planet_units=None, rho_star_units=None,
                 model=None,
                 contact_points_kwargs=None,
                 **kwargs):
        add_citations_to_model(self.__citations__, model=model)

        self.gcc_to_sun = (
            (constants.M_sun / constants.R_sun**3).to(u.g / u.cm**3).value)
        self.G_grav = constants.G.to(u.R_sun**3 / u.M_sun / u.day**2).value

        self.kepler_op = KeplerOp(**kwargs)

        # Parameters
        self.period = tt.as_tensor_variable(period)
        self.t0 = tt.as_tensor_variable(t0)
        self.m_planet = tt.as_tensor_variable(m_planet)
        if m_planet_units is not None:
            self.m_planet *= (1 * m_planet_units).to(u.M_sun).value

        self.a, self.period, self.rho_star, self.r_star, self.m_star = \
            self._get_consistent_inputs(a, period, rho_star, r_star, m_star,
                                        rho_star_units)
        self.m_total = self.m_star + self.m_planet

        self.n = 2 * np.pi / self.period
        self.a_star = self.a * self.m_planet / self.m_total
        self.a_planet = -self.a * self.m_star / self.m_total

        self.K0 = self.n * self.a / self.m_total

        # Set up the contact points calculation
        if contact_points_kwargs is None:
            contact_points_kwargs = dict()

        # Eccentricity
        self.contact_points_op = ContactPointsOp(**contact_points_kwargs)
        if ecc is None:
            self.ecc = None
            self.M0 = 0.5 * np.pi + tt.zeros_like(self.n)
            self.tref = self.t0 - self.M0 / self.n
            incl_factor = 1
        else:
            self.ecc = tt.as_tensor_variable(ecc)
            if omega is None:
                raise ValueError("both e and omega must be provided")
            self.omega = tt.as_tensor_variable(omega)

            self.cos_omega = tt.cos(self.omega)
            self.sin_omega = tt.sin(self.omega)

            opsw = 1 + self.sin_omega
            E0 = 2 * tt.arctan2(tt.sqrt(1-self.ecc)*self.cos_omega,
                                tt.sqrt(1+self.ecc)*opsw)
            self.M0 = E0 - self.ecc * tt.sin(E0)
            self.tref = self.t0 - self.M0 / self.n

            ome2 = 1 - self.ecc**2
            self.K0 /= tt.sqrt(ome2)
            incl_factor = (1 + self.ecc * self.sin_omega) / ome2

        if b is not None:
            if incl is not None or duration is not None:
                raise ValueError("only one of 'incl', 'b', and 'duration' can "
                                 "be given")
            self.b = tt.as_tensor_variable(b)
            self.cos_incl = incl_factor*self.b*self.r_star/self.a_planet
            self.incl = tt.arccos(self.cos_incl)
        elif incl is not None:
            if duration is not None:
                raise ValueError("only one of 'incl', 'b', and 'duration' can "
                                 "be given")
            self.incl = tt.as_tensor_variable(incl)
            self.cos_incl = tt.cos(self.incl)
            self.b = self.a_planet*self.cos_incl/(incl_factor*self.r_star)
        elif duration is not None:
            if self.ecc is None:
                raise ValueError("fitting with duration only works for "
                                 "eccentric orbits")
            self.duration = tt.as_tensor_variable(duration)
            c = tt.sin(np.pi * self.duration * incl_factor / self.period)
            c2 = c * c
            aor = self.a_planet / self.r_star
            esinw = self.ecc * self.sin_omega
            self.b = tt.sqrt((aor**2*c2 - 1)/(c2*esinw**2 + 2*c2*esinw + c2 -
                                              self.ecc**4 + 2*self.ecc**2 - 1))
            self.b *= (1-self.ecc**2)
            self.cos_incl = incl_factor*self.b*self.r_star/self.a_planet
            self.incl = tt.arccos(self.cos_incl)
        else:
            zla = tt.zeros_like(self.a)
            self.incl = 0.5 * np.pi + zla
            self.cos_incl = zla
            self.b = zla

        self.sin_incl = tt.sin(self.incl)
Esempio n. 26
0
                      testval=np.log(34.87846))  # days
    P = pm.Deterministic("P", tt.exp(logP))

    e = pm.Uniform("e", lower=0, upper=1, testval=0.62)
    omega = Angle("omega", testval=80.5 * deg)  # omega_Aa
    Omega = Angle("Omega", testval=110.0 *
                  deg)  # - pi to pi # estimated assuming same as CB disk
    gamma = pm.Uniform("gamma", lower=0, upper=20, testval=10.1)  # km/s

    # uniform on cos incl. testpoint assuming same as CB disk.
    cos_incl = pm.Uniform("cosIncl",
                          lower=0.0,
                          upper=1.0,
                          testval=np.cos(48.0 *
                                         deg))  # radians, 0 to 180 degrees
    incl = pm.Deterministic("incl", tt.arccos(cos_incl))

    # Since we're doing an RV + astrometric fit, M2 now becomes a parameter of the model
    # use a bounded normal to enforce positivity
    PosNormal = pm.Bound(pm.Normal, lower=0.0)
    MAb = PosNormal("MAb", mu=0.3, sd=0.5, testval=0.3)  # solar masses

    t_periastron = pm.Uniform("tPeri",
                              lower=1130.0,
                              upper=1170.0,
                              testval=1159.00)  # + 2400000 days

    orbit = xo.orbits.KeplerianOrbit(
        a=a * au_to_R_sun,
        period=P,
        ecc=e,
Esempio n. 27
0
def make_model(A_re_data,
               A_im_data,
               E_re_data,
               E_im_data,
               Tobs,
               f0,
               fdot,
               fddot,
               sigma,
               hbin,
               lnAlow,
               lnAhigh,
               N,
               start_pt={}):
    f0_mean = f0
    fdot_mean = fdot
    fddot_mean = fddot

    with pm.Model() as model:
        _ = pm.Data('sigma', sigma)
        _ = pm.Data('hbin', hbin)
        _ = pm.Data('Tobs', Tobs)
        _ = pm.Data('N', N)
        A_re_data = pm.Data('A_re_data', A_re_data)
        A_im_data = pm.Data('A_im_data', A_im_data)
        E_re_data = pm.Data('E_re_data', E_re_data)
        E_im_data = pm.Data('E_im_data', E_im_data)

        n_phi = pm.Normal('n_phi',
                          mu=zeros(2),
                          sigma=ones(2),
                          shape=(2, ),
                          testval=start_pt.get('n_phi', randn(2)))
        phi0 = pm.Deterministic('phi0', tt.arctan2(n_phi[1], n_phi[0]))

        dphi_f0 = pm.Normal('dphi_f0', mu=0, sigma=pi, testval=0)
        dphi_fdot = pm.Normal('dphi_fdot', mu=0, sigma=pi, testval=0)
        dphi_fddot = pm.Normal('dphi_fddot', mu=0, sigma=pi, testval=0)

        f0 = pm.Deterministic('f0', f0_mean + dphi_f0 / (2 * pi * Tobs))
        fdot = pm.Deterministic('fdot',
                                fdot_mean + dphi_fdot / (pi * Tobs * Tobs))
        fddot = pm.Deterministic(
            'fddot', fddot_mean + 3.0 * dphi_fddot / (pi * Tobs * Tobs * Tobs))

        cos_iota = pm.Uniform('cos_iota',
                              lower=-1,
                              upper=1,
                              testval=start_pt.get(
                                  'cos_iota', np.random.uniform(low=-1,
                                                                high=1)))
        iota = pm.Deterministic('iota', tt.arccos(cos_iota))

        # This 2-vector gives 2*psi
        n_2psi = pm.Normal('n_2psi',
                           mu=zeros(2),
                           sigma=ones(2),
                           shape=(2, ),
                           testval=start_pt.get('n_2psi', randn(2)))
        psi = pm.Deterministic('psi', tt.arctan2(n_2psi[1], n_2psi[0]) / 2)

        n_ra_dec = pm.Normal('n_ra_dec',
                             mu=zeros(3),
                             sigma=ones(3),
                             shape=(3, ),
                             testval=start_pt.get('nhat', randn(3)))
        nhat = pm.Deterministic(
            'nhat',
            n_ra_dec / pmm.sqrt(tt.tensordot(n_ra_dec, n_ra_dec, axes=1)))
        _ = pm.Deterministic('phi', tt.arctan2(n_ra_dec[1], n_ra_dec[0]))
        _ = pm.Deterministic('theta', tt.arccos(nhat[2]))

        lnA = pm.Uniform('lnA',
                         lower=lnAlow,
                         upper=lnAhigh,
                         testval=start_pt.get(
                             'lnA', np.random.uniform(low=lnAlow,
                                                      high=lnAhigh)))
        A = pm.Deterministic('A', pmm.exp(lnA))

        y_re, y_im = y_fd(Tobs, f0, fdot, fddot, phi0, nhat, cos_iota, psi,
                          hbin, N)
        ((X_re, X_im), (Y_re, Y_im),
         (Z_re, Z_im)) = XYZ_freq(y_re, y_im, Tobs, hbin, N)
        ((A_re, A_im), (E_re, E_im),
         (T_re, T_im)) = AET_XYZ(X_re, X_im, Y_re, Y_im, Z_re, Z_im)

        A_re = pm.Deterministic('A_re', A * A_re)
        A_im = pm.Deterministic('A_im', A * A_im)
        E_re = pm.Deterministic('E_re', A * E_re)
        E_im = pm.Deterministic('E_im', A * E_im)

        snr = pm.Deterministic(
            'SNR',
            tt.sqrt(
                tt.sum(tt.square(A_re / sigma)) +
                tt.sum(tt.square(A_im / sigma)) +
                tt.sum(tt.square(E_re / sigma)) +
                tt.sum(tt.square(E_im / sigma))))

        _ = pm.Normal('A_re_obs', mu=A_re, sigma=sigma, observed=A_re_data)
        _ = pm.Normal('A_im_obs', mu=A_im, sigma=sigma, observed=A_im_data)
        _ = pm.Normal('E_re_obs', mu=E_re, sigma=sigma, observed=E_re_data)
        _ = pm.Normal('E_im_obs', mu=E_im, sigma=sigma, observed=E_im_data)

    return model
Esempio n. 28
0
                            lower=np.log(20),
                            upper=np.log(50.0),
                            testval=np.log(34.88))  # days
    P_inner = pm.Deterministic("PInner", tt.exp(logP_inner))

    e_inner = pm.Uniform("eInner", lower=0, upper=1, testval=0.63)

    omega_inner = Angle("omegaInner", testval=1.415)  # omega_Aa
    Omega_inner = Angle("OmegaInner", testval=1.821)

    # constrained to be i > 90
    cos_incl_inner = pm.Uniform("cosInclInner",
                                lower=-1.0,
                                upper=0.0,
                                testval=-0.659)
    incl_inner = pm.Deterministic("inclInner", tt.arccos(cos_incl_inner))

    MAb = pm.Normal("MAb", mu=0.29, sd=0.5, testval=0.241)  # solar masses

    t_periastron_inner = pm.Uniform("tPeriastronInner",
                                    lower=1140.0,
                                    upper=1180.0,
                                    testval=1159.57)  # + 2400000 + jd0

    orbit_inner = xo.orbits.KeplerianOrbit(
        a=a_inner * au_to_R_sun,
        period=P_inner,
        ecc=e_inner,
        t_periastron=t_periastron_inner,
        omega=omega_inner,
        Omega=Omega_inner,
Esempio n. 29
0
    def evaluate(self, state, dt, motor_signals):

        positions, velocities, rotations = state.positions, state.velocities, state.rotations
        # ALL CONSTRAINTS CAN BE TRANSFORMED TO VELOCITY CONSTRAINTS!
        ##################
        # --- Step 1 --- #
        ##################
        # First, we integrate the applied force F_a acting of each rigid body (like gravity, ...) and
        # we obtain some new velocities newv that tend to violate the constraints.

        # TODO: integrate other forces
        totalforce = np.array(
            [0, 0, 0, 0, 0, 0], dtype='float32'
        )  # total force acting on body outside of constraints

        acceleration = np.array(
            [0, 0, -9.81, 0, 0, 0],
            dtype='float32')  # acceleration of the default frame
        newv = velocities + dt * acceleration[None, None, :]
        originalv = newv

        ##################
        # --- Step 2 --- #
        ##################
        # now enforce the constraints by having corrective impulses

        # convert mass matrices to world coordinates
        # changes every timestep
        M00 = self.lower_inertia_inv
        M01 = T.zeros(shape=(self.batch_size, self.num_moving_bodies, 3, 3))
        M10 = M01
        M11 = T.sum(rotations[:, :, :, None, :, None] *
                    self.upper_inertia_inv[:, :, :, :, None, None] *
                    rotations[:, :, None, :, None, :],
                    axis=(2, 3))
        M0 = T.concatenate([M00, M01], axis=3)
        M1 = T.concatenate([M10, M11], axis=3)
        M = T.concatenate([M0, M1], axis=2)

        # constraints are first dimension! We will need to stack them afterwards!
        J = [
            np.zeros((self.batch_size, 6), dtype="float32")
            for _ in xrange(2 * self.num_constraints)
        ]  # 0 constraints x 0 bodies x 2 objects x 6 velocities
        b_res = [
            np.zeros((self.batch_size, ), dtype="float32")
            for _ in xrange(self.num_constraints)
        ]  # 0 constraints x 0 bodies
        b_error = [
            np.zeros((self.batch_size, ), dtype="float32")
            for _ in xrange(self.num_constraints)
        ]  # 0 constraints x 0 bodies
        C = [
            np.zeros((self.batch_size, ), dtype="float32")
            for _ in xrange(self.num_constraints)
        ]  # 0 constraints x 0 bodies

        c_idx = 0
        for constraint, references, parameters in self.constraints:
            idx1 = references[0]
            idx2 = references[1]
            follows_Newtons_third_law = (idx1 is not None) and (idx2
                                                                is not None)

            # If connected to universe, idx2 is None
            if idx2 is None:
                idx1, idx2 = idx1, idx1
            if idx1 is None:
                idx1, idx2 = idx2, idx2

            if constraint == "ball-and-socket" or constraint == "hinge" or constraint == "fixed":
                r1 = parameters["joint_in_model1_coordinates"][None, :].astype(
                    'float32')
                r2 = parameters["joint_in_model2_coordinates"][None, :]

                r1x = theano_convert_model_to_world_coordinate_no_bias(
                    r1, rotations[:, idx1, :, :])
                r2x = theano_convert_model_to_world_coordinate_no_bias(
                    r2, rotations[:, idx2, :, :])
                ss_r1x = single_skew_symmetric(r1x)
                ss_r2x = single_skew_symmetric(r2x)
                batched_eye = numpy_repeat_new_axis(np.eye(3, dtype='float32'),
                                                    self.batch_size)
                complete_J1 = T.concatenate([-batched_eye, ss_r1x], axis=1)
                complete_J2 = T.concatenate([batched_eye, -ss_r2x], axis=1)
                error = positions[:, idx2, :] + r2x - positions[:,
                                                                idx1, :] - r1x
                for i in xrange(3):
                    J[2 * (c_idx + i) + 0] = complete_J1[:, :, i]
                    J[2 * (c_idx + i) + 1] = complete_J2[:, :, i]

                    b_error[c_idx + i] = error[:, i]
                c_idx += 3

            if constraint == "parallel" or constraint == "plane" or constraint == "slider" or constraint == "fixed":
                batched_eye = numpy_repeat_new_axis(np.eye(3, dtype='float32'),
                                                    self.batch_size)
                batched_zeros = numpy_repeat_new_axis(
                    np.zeros((3, 3), dtype='float32'), self.batch_size)

                complete_J1 = np.concatenate([batched_zeros, -batched_eye],
                                             axis=1)
                complete_J2 = np.concatenate([batched_zeros, batched_eye],
                                             axis=1)

                for i in xrange(3):
                    J[2 * (c_idx + i) + 0] = complete_J1[:, :, i]
                    J[2 * (c_idx + i) + 1] = complete_J2[:, :, i]

                #TODO: this is np?
                #rot_current = theano_dot_last_dimension_matrices(rotations[:,idx2,:,:], rotations[:,idx1,:,:].dimshuffle(0,2,1))
                #rot_diff = np.dot(rot_current, parameters['rot_init'].T)
                #cross = rot_diff.T - rot_diff

                # TODO: add stabilization of this constraint
                b_error[c_idx] = np.zeros(
                    shape=(self.batch_size, ))  #cross[1,2]
                b_error[c_idx + 1] = np.zeros(
                    shape=(self.batch_size, ))  #cross[2,0]
                b_error[c_idx + 2] = np.zeros(
                    shape=(self.batch_size, ))  #cross[0,1]
                c_idx += 3

            if constraint == "plane" or constraint == "slider":
                if follows_Newtons_third_law:
                    n1 = theano_convert_model_to_world_coordinate_no_bias(
                        parameters['axis1_in_model1_coordinates'],
                        rotations[:, idx1, :, :])
                else:
                    n1 = parameters['axis1_in_model1_coordinates']
                    n1 = numpy_repeat_new_axis(n1, self.batch_size)

                complete_J1 = T.concatenate(
                    [-n1, T.zeros(shape=(self.batch_size, 3))], axis=1)
                complete_J2 = -complete_J1

                if follows_Newtons_third_law:
                    J[2 * c_idx + 0] = complete_J1
                J[2 * c_idx + 1] = complete_J2

                if follows_Newtons_third_law:
                    orig_error = theano_convert_model_to_world_coordinate_no_bias(
                        parameters['trans_init_in_model2'],
                        rotations[:, idx2, :, :])
                    pos_error = positions[:,
                                          idx2, :] - positions[:,
                                                               idx1, :] - orig_error
                else:
                    orig_error = parameters['trans_init']
                    orig_error = numpy_repeat_new_axis(orig_error,
                                                       self.batch_size)
                    pos_error = positions[:, idx2, :] - orig_error

                b_error[c_idx] = theano_dot_last_dimension_vectors(
                    pos_error, n1)

                c_idx += 1

            if constraint == "slider":
                if follows_Newtons_third_law:
                    n2 = theano_convert_model_to_world_coordinate_no_bias(
                        parameters['axis2_in_model1_coordinates'],
                        rotations[:, idx1, :, :])
                else:
                    n2 = parameters['axis2_in_model1_coordinates']
                    n2 = numpy_repeat_new_axis(n2, self.batch_size)

                complete_J1 = T.concatenate(
                    [-n2, T.zeros(shape=(self.batch_size, 3))], axis=1)
                complete_J2 = -complete_J1

                if follows_Newtons_third_law:
                    J[2 * c_idx + 0] = complete_J1
                J[2 * c_idx + 1] = complete_J2

                if follows_Newtons_third_law:
                    orig_error = theano_convert_model_to_world_coordinate_no_bias(
                        parameters['trans_init_in_model2'],
                        rotations[:, idx2, :, :])
                    pos_error = positions[:,
                                          idx2, :] - positions[:,
                                                               idx1, :] - orig_error
                else:
                    orig_error = parameters['trans_init']
                    orig_error = numpy_repeat_new_axis(orig_error,
                                                       self.batch_size)
                    pos_error = positions[:, idx2, :] - orig_error

                b_error[c_idx] = theano_dot_last_dimension_vectors(
                    pos_error, n2)

                c_idx += 1

            if constraint == "hinge":
                a2x = theano_convert_model_to_world_coordinate_no_bias(
                    parameters['axis_in_model2_coordinates'][None, :],
                    rotations[:, idx2, :, :])
                b1x = theano_convert_model_to_world_coordinate_no_bias(
                    parameters['axis1_in_model1_coordinates'][None, :],
                    rotations[:, idx1, :, :])
                c1x = theano_convert_model_to_world_coordinate_no_bias(
                    parameters['axis2_in_model1_coordinates'][None, :],
                    rotations[:, idx1, :, :])
                ss_a2x = single_skew_symmetric(a2x)

                batched_zeros = numpy_repeat_new_axis(
                    np.zeros((3, ), dtype='float32'), self.batch_size)

                J[2 * (c_idx + 0) + 0] = T.concatenate([
                    batched_zeros,
                    -theano_dot_last_dimension_vector_matrix(b1x, ss_a2x)
                ],
                                                       axis=1)
                J[2 * (c_idx + 0) + 1] = T.concatenate([
                    batched_zeros,
                    theano_dot_last_dimension_vector_matrix(b1x, ss_a2x)
                ],
                                                       axis=1)
                J[2 * (c_idx + 1) + 0] = T.concatenate([
                    batched_zeros,
                    -theano_dot_last_dimension_vector_matrix(c1x, ss_a2x)
                ],
                                                       axis=1)
                J[2 * (c_idx + 1) + 1] = T.concatenate([
                    batched_zeros,
                    theano_dot_last_dimension_vector_matrix(c1x, ss_a2x)
                ],
                                                       axis=1)

                b_error[c_idx + 0] = theano_dot_last_dimension_vectors(
                    a2x, b1x)
                b_error[c_idx + 1] = theano_dot_last_dimension_vectors(
                    a2x, c1x)
                c_idx += 2

            if constraint == "angular motor":
                ac = parameters['axis_in_model2_coordinates'][None, :]
                a = theano_convert_model_to_world_coordinate_no_bias(
                    ac, rotations[:, idx2, :, :])

                # TODO: remove dimshuffle(0,2,1) by using batched_dot
                rot_current = theano_dot_last_dimension_matrices(
                    rotations[:, idx2, :, :],
                    rotations[:, idx1, :, :].dimshuffle(0, 2, 1))
                rot_init = numpy_repeat_new_axis(parameters['rot_init'].T,
                                                 self.batch_size)
                rot_diff = theano_dot_last_dimension_matrices(
                    rot_current, rot_init)

                traces = rot_diff[:, 0, 0] + rot_diff[:, 1, 1] + rot_diff[:, 2,
                                                                          2]

                # grad when x=-1 or x=1 does not exist for arccos
                theta2 = T.arccos(T.clip(0.5 * (traces - 1), -1 + eps,
                                         1 - eps))
                cross = rot_diff.dimshuffle(0, 2, 1) - rot_diff
                dot2 = cross[:, 1,
                             2] * ac[:,
                                     0] + cross[:, 2,
                                                0] * ac[:,
                                                        1] + cross[:, 0,
                                                                   1] * ac[:,
                                                                           2]

                theta = ((dot2 > 0) * 2 - 1) * theta2

                batched_zeros = numpy_repeat_new_axis(
                    np.zeros((3, ), dtype='float32'), self.batch_size)
                J[2 * c_idx + 0] = T.concatenate([batched_zeros, -a], axis=1)
                J[2 * c_idx + 1] = T.concatenate([batched_zeros, a], axis=1)

                motor_signal = motor_signals[:, parameters["motor_id"]]

                if "min" in parameters and "max" in parameters:
                    motor_min = (parameters["min"] / 180. * np.pi)
                    motor_max = (parameters["max"] / 180. * np.pi)
                    motor_signal = T.clip(motor_signal, motor_min,
                                          motor_max).astype('float32')

                def smallestSignedAngleBetween(x, y):
                    a1 = (x - y) % np.float32(2 * np.pi)
                    b1 = (y - x) % np.float32(2 * np.pi)
                    return T.minimum(a1, b1) * ((a1 > b1) * 2 - 1)

                error_signal = -smallestSignedAngleBetween(theta, motor_signal)

                if parameters["type"] == "velocity":
                    b_error[c_idx] = motor_signal
                elif parameters["type"] == "position":
                    if "delta" in parameters and "motor_velocity" in parameters:
                        velocity = parameters["motor_velocity"] / 180. * np.pi
                        b_error[c_idx] = dt * T.clip(
                            (abs(error_signal) > parameters["delta"]) *
                            error_signal * parameters["motor_gain"], -velocity,
                            velocity)
                    else:
                        b_error[c_idx] = dt * error_signal * parameters[
                            "motor_gain"]

                c_idx += 1

            if constraint == "linear motor":

                ac = parameters['axis_in_model2_coordinates'][None, :]
                a = theano_convert_model_to_world_coordinate_no_bias(
                    ac, rotations[:, idx2, :, :])

                batched_zeros = numpy_repeat_new_axis(
                    np.zeros((3, ), dtype='float32'), self.batch_size)

                if follows_Newtons_third_law:
                    J[2 * c_idx + 0] = T.concatenate([-a, batched_zeros],
                                                     axis=1)
                J[2 * c_idx + 1] = T.concatenate([a, batched_zeros], axis=1)

                if follows_Newtons_third_law:
                    position = positions[:,
                                         idx2, :] - positions[:,
                                                              idx1, :] - parameters[
                                                                  'pos_init']
                else:
                    position = positions[:, idx2, :] - parameters['pos_init']

                motor_signal = motor_signals[:, parameters["motor_id"]]
                if "min" in parameters and "max" in parameters:
                    motor_min = parameters["min"]
                    motor_max = parameters["max"]
                    motor_signal = T.clip(motor_signal, motor_min, motor_max)

                error_signal = theano_dot_last_dimension_vectors(
                    position, a) - motor_signal

                if parameters["servo"] == "velocity":
                    b_error[c_idx] = motor_signal
                elif parameters["servo"] == "position":
                    if "delta" in parameters and "motor_velocity" in parameters:
                        velocity = parameters["motor_velocity"]
                        b_error[c_idx] = dt * T.clip(
                            (abs(error_signal) > parameters["delta"]) *
                            error_signal * parameters["motor_gain"], -velocity,
                            velocity)
                    else:
                        b_error[c_idx] = dt * error_signal * parameters[
                            "motor_gain"]

                c_idx += 1

            if constraint == "angular limit":
                angle = parameters["angle"] / 180. * np.pi

                ac = parameters['axis_in_model2_coordinates'][None, :]
                a = theano_convert_model_to_world_coordinate_no_bias(
                    parameters['axis_in_model2_coordinates'],
                    rotations[:, idx2, :, :])

                rot_current = theano_dot_last_dimension_matrices(
                    rotations[:, idx2, :, :],
                    rotations[:, idx1, :, :].dimshuffle(0, 2, 1))
                rot_init = parameters['rot_init'].T
                rot_diff = theano_dot_last_dimension_matrices(
                    rot_current, rot_init)

                traces = rot_diff[:, 0, 0] + rot_diff[:, 1, 1] + rot_diff[:, 2,
                                                                          2]

                theta2 = T.arccos(T.clip(0.5 * (traces - 1), -1 + eps,
                                         1 - eps))
                cross = rot_diff.dimshuffle(0, 2, 1) - rot_diff
                dot2 = cross[:, 1,
                             2] * ac[:,
                                     0] + cross[:, 2,
                                                0] * ac[:,
                                                        1] + cross[:, 0,
                                                                   1] * ac[:,
                                                                           2]

                theta = ((dot2 > 0) * 2 - 1) * theta2

                batched_zeros = numpy_repeat_new_axis(
                    np.zeros((3, ), dtype='float32'), self.batch_size)

                if parameters["angle"] < 0:
                    if follows_Newtons_third_law:
                        J[2 * c_idx + 0] = T.concatenate([batched_zeros, -a])
                    J[2 * c_idx + 1] = T.concatenate([batched_zeros, a])
                else:
                    if follows_Newtons_third_law:
                        J[2 * c_idx + 0] = T.concatenate([batched_zeros, a])
                    J[2 * c_idx + 1] = T.concatenate([batched_zeros, -a])

                b_error[c_idx] = T.abs_(angle - theta)

                if parameters["angle"] > 0:
                    b_error[c_idx] = angle - theta
                    self.C[c_idx] = (theta > angle)
                else:
                    b_error[c_idx] = theta - angle
                    self.C[c_idx] = (theta < angle)

                c_idx += 1

            if constraint == "linear limit":

                offset = parameters["offset"]
                if follows_Newtons_third_law:
                    ac = parameters['axis_in_model1_coordinates'][None, :]
                    a = theano_convert_model_to_world_coordinate_no_bias(
                        ac, rotations[:, idx1, :, :])
                else:
                    a = parameters['axis_in_model1_coordinates'][None, :]

                if offset < 0:
                    if follows_Newtons_third_law:
                        J[2 * c_idx + 0] = T.concatenate([-a, batched_zeros])
                    J[2 * c_idx + 1] = T.concatenate([a, batched_zeros])
                else:
                    if follows_Newtons_third_law:
                        J[2 * c_idx + 0] = T.concatenate([a, batched_zeros])
                    J[2 * c_idx + 1] = T.concatenate([-a, batched_zeros])

                if follows_Newtons_third_law:
                    position = positions[:,
                                         idx2, :] - positions[:,
                                                              idx1, :] - parameters[
                                                                  'pos_init']
                else:
                    position = positions[:, idx2, :] - parameters['pos_init']

                current_offset = theano_dot_last_dimension_vectors(position, a)

                if parameters["offset"] > 0:
                    b_error[c_idx] = offset - current_offset
                    self.C[c_idx] = (current_offset > offset)
                else:
                    b_error[c_idx] = offset - current_offset
                    self.C[c_idx] = (current_offset < offset)

                c_idx += 1

            if constraint == "ground":
                r = self.radii[idx1].astype('float32')
                J[2 * c_idx + 0] = numpy_repeat_new_axis(
                    np.array([0, 0, 1, 0, 0, 0], dtype='float32'),
                    self.batch_size)
                J[2 * c_idx + 1] = numpy_repeat_new_axis(
                    np.array([0, 0, 0, 0, 0, 0], dtype='float32'),
                    self.batch_size)

                # TODO: use T.maximum
                b_error[c_idx] = T.clip(
                    positions[:, idx1, Z] - r + parameters["delta"],
                    np.finfo('float32').min, 0)
                b_res[c_idx] = parameters["alpha"] * newv[:, idx1, Z]
                C[c_idx] = positions[:, idx1, Z] - r
                c_idx += 1

            if constraint == "ground" and parameters["mu"] != 0:
                r = self.radii[idx1].astype('float32')
                for i in xrange(2):
                    if i == 0:
                        J[2 * (c_idx + i) + 0] = numpy_repeat_new_axis(
                            np.array([0, 1, 0, -r, 0, 0], dtype='float32'),
                            self.batch_size)
                    else:
                        J[2 * (c_idx + i) + 0] = numpy_repeat_new_axis(
                            np.array([1, 0, 0, 0, r, 0], dtype='float32'),
                            self.batch_size)
                    J[2 * (c_idx + i) + 1] = numpy_repeat_new_axis(
                        np.array([0, 0, 0, 0, 0, 0], dtype='float32'),
                        self.batch_size)
                    C[c_idx + i] = positions[:, idx1, Z] - r
                c_idx += 2

            if constraint == "ground" and parameters[
                    "torsional_friction"] and parameters["mu"] != 0:
                r = self.radii[idx1].astype('float32')
                J[2 * c_idx + 0] = numpy_repeat_new_axis(
                    np.array([0, 0, 0, 0, 0, r], dtype='float32'),
                    self.batch_size)
                J[2 * c_idx + 1] = numpy_repeat_new_axis(
                    np.array([0, 0, 0, 0, 0, 0], dtype='float32'),
                    self.batch_size)
                C[c_idx] = positions[:, idx1, Z] - r
                c_idx += 1

        zipped_indices = [
            j for i in zip(self.zero_index, self.one_index) for j in i
        ]

        mass_matrix = M[:, zipped_indices, :, :].reshape(
            shape=(self.batch_size, self.num_constraints, 2, 6, 6))

        J = theano_stack_batched_integers_mixed_numpy(
            J, expected_shape=(self.num_constraints, 2, self.batch_size,
                               6)).dimshuffle(2, 0, 1, 3)

        C = theano_stack_batched_integers_mixed_numpy(
            C, expected_shape=(self.num_constraints,
                               self.batch_size)).dimshuffle(1, 0)
        b_res = theano_stack_batched_integers_mixed_numpy(
            b_res, expected_shape=(self.num_constraints,
                                   self.batch_size)).dimshuffle(1, 0)
        b_error = theano_stack_batched_integers_mixed_numpy(
            b_error, expected_shape=(self.num_constraints,
                                     self.batch_size)).dimshuffle(1, 0)

        self.impulses_P = self.warm_start * self.impulses_P

        # TODO: batch-dot-product
        m_eff = 1. / T.sum(
            T.sum(J[:, :, :, None, :] * mass_matrix, axis=4) * J, axis=(2, 3))

        k = m_eff * (self.w**2)
        c = m_eff * 2 * self.zeta * self.w

        CFM = 1. / (c + dt * k)
        ERP = dt * k / (c + dt * k)
        m_c = 1. / (1. / m_eff + CFM)
        b = ERP / dt * b_error + b_res

        for iteration in xrange(self.projected_gauss_seidel_iterations):
            # this changes every iteration
            v = newv[:, zipped_indices, :].reshape(shape=(self.batch_size,
                                                          self.num_constraints,
                                                          2, 6))

            lamb = -m_c * (T.sum(J * v, axis=(2, 3)) + CFM * self.impulses_P +
                           b)
            self.impulses_P += lamb

            if self.do_impulse_clipping:
                if self.do_friction_clipping:
                    clipping_force = self.impulses_P[:, self.clipping_idx]
                    clipping_limit = abs(self.clipping_a * clipping_force +
                                         self.clipping_b * dt)
                else:
                    clipping_limit = self.clipping_b * dt
                self.impulses_P = T.clip(self.impulses_P, -clipping_limit,
                                         clipping_limit)

            if self.has_conditional_constraints:
                applicable = (1.0 * (C <= 0)) * (1 - (self.only_when_positive *
                                                      (self.impulses_P <= 0)))
                self.impulses_P = self.impulses_P * applicable

            # TODO: batch-dot-product
            result = T.sum(mass_matrix * J[:, :, :, None, :],
                           axis=4) * self.impulses_P[:, :, None, None]
            result = result.reshape(
                (self.batch_size, 2 * self.num_constraints, 6))

            r = []
            for constraint in self.map_object_to_constraint:
                idx_list = np.array(constraint,
                                    dtype='int64')  # deal with empty lists
                delta_v = T.sum(result[:, idx_list, :], axis=1)
                r.append(delta_v)
            newv = newv + T.stack(r, axis=1)
        #print
        #print theano.printing.debugprint(newv)
        return newv
def _get_compiled_theano_functions():
    # Planet masses: m1,m2
    m1, m2 = T.dscalars(2)
    mstar = 1
    mu1 = m1 * mstar / (mstar + m1)
    mu2 = m2 * mstar / (mstar + m2)
    eta1 = mstar + m1
    eta2 = mstar + m2
    beta1 = mu1 * T.sqrt(eta1 / mstar) / (mu1 + mu2)
    beta2 = mu2 * T.sqrt(eta2 / mstar) / (mu1 + mu2)
    j, k = T.lscalars('jk')
    s = (j - k) / k

    # Dynamical variables:
    dyvars = T.vector()
    s1, s2, psi, phi, Omega, I1, I2, Psi, Phi, Rtilde = [
        dyvars[i] for i in range(10)
    ]
    l1 = phi - 0.5 * k * psi
    l2 = phi + 0.5 * k * psi
    gamma1 = s1 - (1 + s) * l2 + s * l1
    gamma2 = s2 - (1 + s) * l2 + s * l1
    Gamma1 = I1
    Gamma2 = I2
    L1 = Phi / 2 - Psi / k - s * (I1 + I2)
    L2 = Phi / 2 + Psi / k + (s + 1) * (I1 + I2)
    Cz = -1 * Rtilde

    R = L1 + L2 - Gamma1 - Gamma2 - Cz
    G1 = L1 - Gamma1
    G2 = L2 - Gamma2

    r2_by_r1 = (L2 - L1 - Gamma2 + Gamma1) / (L1 + L2 - Gamma1 - Gamma2 - R)
    rho1 = 0.5 * R * (1 + r2_by_r1)
    rho2 = 0.5 * R * (1 - r2_by_r1)

    a1 = (L1 / beta1)**2
    e1 = T.sqrt(1 - (1 - (Gamma1 / L1))**2)

    a2 = (L2 / beta2)**2
    e2 = T.sqrt(1 - (1 - (Gamma2 / L2))**2)

    cos_inc1 = 1 - rho1 / G1
    cos_inc2 = 1 - rho2 / G2
    inc1 = T.arccos(cos_inc1)
    inc2 = T.arccos(cos_inc2)

    l1_r = l1 - Omega
    l2_r = l2 - Omega

    Omega1_r = T.constant(np.pi / 2) - Omega
    Omega2_r = Omega1_r - T.constant(np.pi)

    pomega1 = -1 * gamma1
    pomega2 = -1 * gamma2

    pomega1_r = pomega1 - Omega
    pomega2_r = pomega2 - Omega

    omega1 = pomega1_r - Omega1_r
    omega2 = pomega2_r - Omega2_r

    Hkep = -0.5 * T.sqrt(eta1) * beta1 / a1 - 0.5 * T.sqrt(eta2) * beta2 / a2

    ko = KeplerOp()
    M1 = l1_r - pomega1_r
    M2 = l2_r - pomega2_r
    sinf1, cosf1 = ko(M1, e1 + T.zeros_like(M1))
    sinf2, cosf2 = ko(M2, e2 + T.zeros_like(M2))
    #
    n1 = T.sqrt(eta1 / mstar) * a1**(-3 / 2)
    n2 = T.sqrt(eta2 / mstar) * a2**(-3 / 2)
    Hint_dir, Hint_ind, r1, r2, v1, v2 = calc_Hint_components_sinf_cosf(
        a1, a2, e1, e2, inc1, inc2, omega1, omega2, Omega1_r, Omega2_r, n1, n2,
        sinf1, cosf1, sinf2, cosf2)
    eps = m1 * m2 / (mu1 + mu2) / T.sqrt(mstar)
    Hpert = (Hint_dir + Hint_ind / mstar)
    Htot = Hkep + eps * Hpert

    #####################################################
    # Set parameters for compiling functions with Theano
    #####################################################

    # 'ins' will set the inputs of Theano functions compiled below
    #   Note: 'extra_ins' will be passed as values of object attributes
    #   of the 'ResonanceEquations' class 'defined below
    extra_ins = [m1, m2, j, k]
    givens = []
    ins = [dyvars] + extra_ins
    orbels = [
        a1, e1, inc1, l1_r, pomega1_r, Omega1_r, a2, e2, inc2, l2_r, pomega2_r,
        Omega2_r
    ]
    #  Conservative flow
    gradHtot = T.grad(Htot, wrt=dyvars)
    hessHtot = theano.gradient.hessian(Htot, wrt=dyvars)
    Jtens = T.as_tensor(_get_Omega_matrix(5))
    H_flow_vec = Jtens.dot(gradHtot)
    H_flow_jac = Jtens.dot(hessHtot)

    ##########################
    # Compile Theano functions
    ##########################
    orbels_fn = theano.function(inputs=ins,
                                outputs=orbels,
                                givens=givens,
                                on_unused_input='ignore')

    rv1_fn = theano.function(inputs=ins,
                             outputs=r1 + v1,
                             givens=givens,
                             on_unused_input='ignore')
    rv2_fn = theano.function(inputs=ins,
                             outputs=r2 + v2,
                             givens=givens,
                             on_unused_input='ignore')

    Htot_fn = theano.function(inputs=ins,
                              outputs=Htot,
                              givens=givens,
                              on_unused_input='ignore')

    Hpert_fn = theano.function(inputs=ins,
                               outputs=Hpert,
                               givens=givens,
                               on_unused_input='ignore')

    Hpert_components_fn = theano.function(inputs=ins,
                                          outputs=[Hint_dir, Hint_ind],
                                          givens=givens,
                                          on_unused_input='ignore')

    H_flow_vec_fn = theano.function(inputs=ins,
                                    outputs=H_flow_vec,
                                    givens=givens,
                                    on_unused_input='ignore')

    H_flow_jac_fn = theano.function(inputs=ins,
                                    outputs=H_flow_jac,
                                    givens=givens,
                                    on_unused_input='ignore')
    return dict({
        'orbital_elements': orbels_fn,
        'Hamiltonian': Htot_fn,
        'Hpert': Hpert_fn,
        'Hpert_components': Hpert_components_fn,
        'Hamiltonian_flow': H_flow_vec_fn,
        'Hamiltonian_flow_jacobian': H_flow_jac_fn,
        'positions_and_velocities1': rv1_fn,
        'positions_and_velocities2': rv2_fn
    })
Esempio n. 31
0
 def dist(self, X, Y):
     XY = tensor.sum((X * Y), axis=0)
     XY[XY > 1] = 1
     U = tensor.arccos(XY)
     return U.norm()
Esempio n. 32
0
    def evaluate(self, dt, positions, velocities, rot_matrices, motor_signals):

        # ALL CONSTRAINTS CAN BE TRANSFORMED TO VELOCITY CONSTRAINTS!
        ##################
        # --- Step 1 --- #
        ##################
        # First, we integrate the applied force F_a acting of each rigid body (like gravity, ...) and
        # we obtain some new velocities v2' that tends to violate the constraints.

        totalforce = np.array(
            [0, 0, 0, 0, 0, 0], dtype='float32'
        )  # total force acting on body outside of constraints
        acceleration = np.array(
            [0, 0, -9.81, 0, 0, 0],
            dtype='float32')  # acceleration of the default frame
        newv = velocities + dt * acceleration[None, :]
        originalv = newv

        ##################
        # --- Step 2 --- #
        ##################
        # now enforce the constraints by having corrective impulses
        # convert mass matrices to world coordinates
        M = T.zeros(shape=(self.massMatrices.shape[0], 6, 6))

        M00 = self.lower_inertia_inv
        M01 = T.zeros(shape=(self.massMatrices.shape[0], 3, 3))
        M10 = M01
        M11 = T.sum(rot_matrices.dimshuffle(0, 1, 'x', 2, 'x') *
                    self.upper_inertia_inv.dimshuffle(0, 1, 2, 'x', 'x') *
                    rot_matrices.dimshuffle(0, 'x', 1, 'x', 2),
                    axis=(1, 2))
        M0 = T.concatenate([M00, M01], axis=2)
        M1 = T.concatenate([M10, M11], axis=2)
        M = T.concatenate([M0, M1], axis=1)

        #self.P = np.zeros((self.num_constraints,))# constant
        # changes every timestep
        J = [np.zeros((1, 6)) for _ in xrange(2 * self.num_constraints)
             ]  # 0 constraints x 2 objects x 6 states
        b_res = [0 for _ in xrange(self.num_constraints)]  # 0 constraints
        b_error = [0 for _ in xrange(self.num_constraints)]  # 0 constraints
        C = [0 for _ in xrange(self.num_constraints)]  # 0 constraints

        c_idx = 0
        for constraint, references, parameters in self.constraints:
            idx1 = references[0]
            if references[1] is not None:
                idx2 = references[1]
            else:
                idx2 = None

            if constraint == "ball-and-socket" or constraint == "hinge" or constraint == "fixed":
                r1x = theano_convert_model_to_world_coordinate_no_bias(
                    parameters["joint_in_model1_coordinates"],
                    rot_matrices[idx1, :, :])
                r2x = theano_convert_model_to_world_coordinate_no_bias(
                    parameters["joint_in_model2_coordinates"],
                    rot_matrices[idx2, :, :])
                ss_r1x = single_skew_symmetric(r1x)
                ss_r2x = single_skew_symmetric(r2x)
                complete_J1 = T.concatenate(
                    [-np.eye(3, dtype='float32'), ss_r1x])
                complete_J2 = T.concatenate(
                    [np.eye(3, dtype='float32'), -ss_r2x])
                error = positions[idx2, :] + r2x - positions[idx1, :] - r1x
                for i in xrange(3):
                    J[2 * (c_idx + i) + 0] = complete_J1[:, i]
                    J[2 * (c_idx + i) + 1] = complete_J2[:, i]

                    b_error[c_idx + i] = error[i]
                c_idx += 3

            if constraint == "slider" or constraint == "fixed":
                complete_J1 = np.concatenate([np.zeros((3, 3)),
                                              -np.eye(3)]).astype('float32')
                complete_J2 = np.concatenate([np.zeros((3, 3)),
                                              np.eye(3)]).astype('float32')

                for i in xrange(3):
                    J[2 * (c_idx + i) + 0] = complete_J1[:, i]
                    J[2 * (c_idx + i) + 1] = complete_J2[:, i]

                rot_current = np.dot(rot_matrices[idx2, :, :],
                                     rot_matrices[idx1, :, :].T)
                rot_diff = np.dot(rot_current, parameters['rot_init'].T)
                cross = rot_diff.T - rot_diff
                # TODO: add stabilization of this constraint
                b_error[c_idx] = 0  #cross[1,2]
                b_error[c_idx + 1] = 0  #cross[2,0]
                b_error[c_idx + 2] = 0  #cross[0,1]
                c_idx += 3

            if constraint == "hinge":
                a2x = theano_convert_model_to_world_coordinate_no_bias(
                    parameters['axis_in_model2_coordinates'],
                    rot_matrices[idx2, :, :])
                b1x = theano_convert_model_to_world_coordinate_no_bias(
                    parameters['axis1_in_model1_coordinates'],
                    rot_matrices[idx1, :, :])
                c1x = theano_convert_model_to_world_coordinate_no_bias(
                    parameters['axis2_in_model1_coordinates'],
                    rot_matrices[idx1, :, :])
                ss_a2x = single_skew_symmetric(a2x)

                J[2 * (c_idx + 0) + 0] = T.concatenate(
                    [np.zeros((3, ), dtype='float32'), -T.dot(b1x, ss_a2x)])
                J[2 * (c_idx + 0) + 1] = T.concatenate(
                    [np.zeros((3, ), dtype='float32'),
                     T.dot(b1x, ss_a2x)])
                J[2 * (c_idx + 1) + 0] = T.concatenate(
                    [np.zeros((3, ), dtype='float32'), -T.dot(c1x, ss_a2x)])
                J[2 * (c_idx + 1) + 1] = T.concatenate(
                    [np.zeros((3, ), dtype='float32'),
                     T.dot(c1x, ss_a2x)])

                b_error[c_idx + 0] = T.sum(a2x * b1x)
                b_error[c_idx + 1] = T.sum(a2x * c1x)
                c_idx += 2

            if constraint == "limit":
                angle = parameters["angle"] / 180. * np.pi
                a = theano_convert_model_to_world_coordinate_no_bias(
                    parameters['axis_in_model1_coordinates'],
                    rot_matrices[idx1, :, :])
                rot_current = np.dot(rot_matrices[idx2, :, :],
                                     rot_matrices[idx1, :, :].T)
                rot_diff = np.dot(rot_current,
                                  quat_to_rot_matrix(parameters['rot_init']).T)
                theta2 = np.arccos(0.5 * (np.trace(rot_diff) - 1))
                cross = rot_diff.T - rot_diff
                dot2 = cross[1, 2] * a[0] + cross[2, 0] * a[1] + cross[
                    0, 1] * a[2]
                theta = ((dot2 > 0) * 2 - 1) * theta2

                if parameters["angle"] < 0:
                    J[2 * c_idx + 0] = np.concatenate(
                        [np.zeros((3, ), dtype='float32'), -a])
                    J[2 * c_idx + 1] = np.concatenate(
                        [np.zeros((3, ), dtype='float32'), a])
                else:
                    J[2 * c_idx + 0] = np.concatenate(
                        [np.zeros((3, ), dtype='float32'), a])
                    J[2 * c_idx + 1] = np.concatenate(
                        [np.zeros((3, ), dtype='float32'), -a])

                b_error[c_idx] = np.abs(angle - theta)
                if parameters["angle"] > 0:
                    b_error[c_idx] = angle - theta
                    C[c_idx] = (theta > angle)
                else:
                    b_error[c_idx] = theta - angle
                    C[c_idx] = (theta < angle)

                c_idx += 1

            if constraint == "motor":
                a = theano_convert_model_to_world_coordinate_no_bias(
                    parameters['axis_in_model1_coordinates'],
                    rot_matrices[idx1, :, :])

                rot_current = T.dot(rot_matrices[idx2, :, :],
                                    rot_matrices[idx1, :, :].T)
                rot_diff = T.dot(rot_current, parameters['rot_init'].T)
                theta2 = T.arccos(
                    T.clip(0.5 * (T.nlinalg.trace(rot_diff) - 1), -1 + eps,
                           1 - eps))
                cross = rot_diff.T - rot_diff
                dot2 = cross[1, 2] * a[0] + cross[2, 0] * a[1] + cross[
                    0, 1] * a[2]

                theta = ((dot2 > 0) * 2 - 1) * theta2

                J[2 * c_idx + 0] = T.concatenate(
                    [np.zeros((3, ), dtype='float32'), -a])
                J[2 * c_idx + 1] = T.concatenate(
                    [np.zeros((3, ), dtype='float32'), a])

                motor_signal = motor_signals[parameters["motor_id"]]
                motor_min = parameters["min"] / 180. * np.pi
                motor_max = parameters["max"] / 180. * np.pi

                motor_signal = T.clip(motor_signal, motor_min, motor_max)

                if parameters["type"] == "velocity":
                    b_error[c_idx] = motor_signal
                elif parameters["type"] == "position":
                    if "delta" in parameters and parameters["delta"] > 0:
                        b_error[c_idx] = dt * (
                            abs(theta - motor_signal) > parameters["delta"]
                        ) * (theta -
                             motor_signal) * parameters["motor_velocity"]
                    else:
                        b_error[c_idx] = dt * (theta - motor_signal
                                               ) * parameters["motor_velocity"]

                #print c_idx

                c_idx += 1

            if constraint == "ground":
                r = self.radii[idx1].astype('float32')
                J[2 * c_idx + 0] = np.array([0, 0, 1, 0, 0, 0],
                                            dtype='float32')
                J[2 * c_idx + 1] = np.array([0, 0, 0, 0, 0, 0],
                                            dtype='float32')

                b_error[c_idx] = T.clip(
                    positions[idx1, Z] - r + parameters["delta"],
                    np.finfo('float32').min, 0)
                b_res[c_idx] = parameters["alpha"] * newv[idx1, Z]
                C[c_idx] = positions[idx1, Z] - r
                c_idx += 1

            if constraint == "ground" and parameters["mu"] != 0:
                r = self.radii[idx1].astype('float32')
                for i in xrange(2):
                    if i == 0:
                        J[2 * (c_idx + i) + 0] = np.array([0, 1, 0, -r, 0, 0],
                                                          dtype='float32')
                    else:
                        J[2 * (c_idx + i) + 0] = np.array([1, 0, 0, 0, r, 0],
                                                          dtype='float32')
                    J[2 * (c_idx + i) + 1] = np.array([0, 0, 0, 0, 0, 0],
                                                      dtype='float32')
                    C[c_idx + i] = positions[idx1, Z] - r
                c_idx += 2

            if constraint == "ground" and parameters[
                    "torsional_friction"] and parameters["mu"] != 0:
                r = self.radii[idx1].astype('float32')
                J[2 * c_idx + 0] = np.array([0, 0, 0, 0, 0, r],
                                            dtype='float32')
                J[2 * c_idx + 1] = np.array([0, 0, 0, 0, 0, 0],
                                            dtype='float32')
                C[c_idx] = positions[idx1, Z] - r
                c_idx += 1

        mass_matrix = T.concatenate(
            (T.stack([M[i, None, :, :] for i in self.zero_index], axis=0),
             T.stack([M[i, None, :, :] for i in self.one_index], axis=0)),
            axis=1)

        J = T.stack(J, axis=0).reshape(shape=(self.num_constraints, 2, 6))
        C = theano_stack_integers_mixed_numpy(C)
        #v = np.zeros((self.num_constraints,2,6))  # 0 constraints x 2 objects x 6 states
        b_res = theano_stack_integers_mixed_numpy(b_res)
        b_error = theano_stack_integers_mixed_numpy(b_error)

        for iteration in xrange(self.num_iterations):
            # changes every iteration
            v = T.concatenate(
                (T.stack([newv[i, None, :] for i in self.zero_index], axis=0),
                 T.stack([newv[j, None, :] for j in self.one_index], axis=0)),
                axis=1)

            m_eff = 1. / T.sum(
                T.sum(J[:, :, None, :] * mass_matrix, axis=3) * J, axis=(1, 2))

            k = m_eff * (self.w**2)
            c = m_eff * 2 * self.zeta * self.w

            CFM = 1. / (c + dt * k)
            ERP = dt * k / (c + dt * k)

            m_c = 1. / (1. / m_eff + CFM)

            b = ERP / dt * b_error + b_res

            #theano_to_print.extend([b[i] for i in [3,7,11,15]])

            lamb = -m_c * (T.sum(J * v, axis=(1, 2)) + CFM * self.P + b)

            self.P += lamb

            #print J[[39,61,65,69],:]
            #print np.sum(lamb**2), np.sum(self.P**2)
            clipping_force = T.concatenate(
                [self.P[j].dimshuffle('x') for j in self.clipping_idx], axis=0)
            clipping_limit = abs(self.clipping_a * clipping_force +
                                 self.clipping_b * dt)
            self.P = T.clip(self.P, -clipping_limit, clipping_limit)
            applicable = (1 - (self.only_when_positive *
                               (1 - (self.P >= 0)))) * (C <= 0)
            result = T.sum(mass_matrix * J[:, :, None, :],
                           axis=3) * self.P[:, None,
                                            None] * applicable[:, None, None]
            #theano_to_print.extend([T.sum(abs(result[i]),axis=(-1,-2)) for i in [3,7,11,15]])
            result = result.reshape((2 * self.num_constraints, 6))
            #theano_to_print.extend([T.sum(abs(result[i]),axis=(-1)) for i in [6,14,22,30]])

            r = []
            for i in xrange(len(self.map_object_to_constraint)):
                delta_v = T.sum(T.stack(
                    [result[j, :] for j in self.map_object_to_constraint[i]],
                    axis=0),
                                axis=0)
                r.append(delta_v)

            #theano_to_print.extend([T.sum(abs(newv[i]),axis=(-1)) for i in [4,3,5,11]])
            delta_v = T.stack(r, axis=0)
            #newv = originalv + delta_v
            newv = newv + delta_v

        #print

        return newv
Esempio n. 33
0
    def __init__(self,
                 numpy_rng,
                 theano_rng=None,
                 n_ins=40 * 3,
                 layers_types=[ReLU, ReLU, ReLU, ReLU, ReLU],
                 layers_sizes=[1024, 1024, 1024, 1024],
                 n_outs=100,
                 n_embs=2,
                 loss='cos2',
                 rho=0.9,
                 eps=1.E-6,
                 max_norm=0.,
                 debugprint=False):
        self.layers = []
        self.params = []
        self.n_layers = len(layers_types)
        self.layers_types = layers_types
        assert self.n_layers > 0
        self.max_norm = max_norm
        self._rho = rho  # ``momentum'' for adadelta
        self._eps = eps  # epsilon for adadelta
        self._accugrads = []  # for adadelta
        self._accudeltas = []  # for adadelta

        if theano_rng == None:
            theano_rng = RandomStreams(numpy_rng.randint(2**30))

        self.x1 = T.fmatrix('x1')
        self.x2 = T.fmatrix('x2')
        self.x3 = T.fmatrix('x3')
        self.y12s = [T.ivector('y12') for _ in n_embs]
        self.y13s = [T.ivector('y13') for _ in n_embs]
        self.n_embs = n_embs

        self.layers_ins = [n_ins] + layers_sizes[:-1]
        self.layers_outs = layers_sizes
        layer_input1 = self.x1
        layer_input2 = self.x2
        layer_input3 = self.x3
        layer_embs1 = [None for _ in n_embs]
        layer_embs2 = [None for _ in n_embs]
        layer_embs3 = [None for _ in n_embs]

        for layer_ind, (layer_type, n_in, n_out) in enumerate(
                zip(layers_types[:-1], self.layers_ins, self.layers_outs)):
            this_layer1 = layer_type(rng=numpy_rng,
                                     input=layer_input1,
                                     n_in=n_in,
                                     n_out=n_out)  #, cap=6.)
            assert hasattr(this_layer1, 'output')
            layer_input1 = this_layer1.output
            self.params.extend(this_layer1.params)
            self._accugrads.extend([
                build_shared_zeros(t.shape.eval(), 'accugrad')
                for t in this_layer1.params
            ])
            self._accudeltas.extend([
                build_shared_zeros(t.shape.eval(), 'accudelta')
                for t in this_layer1.params
            ])
            self.layers.append(this_layer1)
            this_layer2 = layer_type(rng=numpy_rng,
                                     input=layer_input2,
                                     n_in=n_in,
                                     n_out=n_out,
                                     W=this_layer1.W,
                                     b=this_layer1.b)  #, cap=6.)
            assert hasattr(this_layer2, 'output')
            this_layer3 = layer_type(rng=numpy_rng,
                                     input=layer_input2,
                                     n_in=n_in,
                                     n_out=n_out,
                                     W=this_layer1.W,
                                     b=this_layer1.b)  #, cap=6.)
            assert hasattr(this_layer3, 'output')
            layer_input2 = this_layer2.output
            self.layers.append(this_layer2)
            layer_input3 = this_layer3.output
            self.layers.append(this_layer3)

        for i_emb in xrange(n_embs):
            emb_layer_type = layers_types[-1]
            n_in = layers_sizes[-1]
            n_out = n_outs
            this_emb_layer1 = emb_layer_type(rng=numpy_rng,
                                             input=layer_input1,
                                             n_in=n_in,
                                             n_out=n_out)  #, cap=6.)
            assert hasattr(this_emb_layer1, 'output')
            layer_embs1[i_emb] = this_emb_layer1.output
            self.params.extend(this_emb_layer1.params)
            self._accugrads.extend([
                build_shared_zeros(t.shape.eval(), 'accugrad')
                for t in this_emb_layer1.params
            ])
            self._accudeltas.extend([
                build_shared_zeros(t.shape.eval(), 'accudelta')
                for t in this_emb_layer1.params
            ])
            this_emb_layer2 = emb_layer_type(rng=numpy_rng,
                                             input=layer_input2,
                                             n_in=n_in,
                                             n_out=n_out,
                                             W=this_emb_layer1.W,
                                             b=this_emb_layer1.b)  #, cap=6.)
            assert hasattr(this_emb_layer2, 'output')
            layer_embs2[i_emb] = this_emb_layer2.output
            this_emb_layer3 = emb_layer_type(rng=numpy_rng,
                                             input=layer_input3,
                                             n_in=n_in,
                                             n_out=n_out,
                                             W=this_emb_layer1.W,
                                             b=this_emb_layer1.b)  #, cap=6.)
            assert hasattr(this_emb_layer3, 'output')
            layer_embs3[i_emb] = this_emb_layer3.output
            self.layers.append(this_emb_layer1)
            self.layers.append(this_emb_layer2)
            self.layers.append(this_emb_layer3)

        L2 = 0.
        for param in self.params:
            L2 += T.sum(param**2)
        L1 = 0.
        for param in self.params:
            L1 += T.sum(abs(param))

        self.cos2_costs = []
        self.acos_costs = []

        for i_emb in xrange(n_embs):
            y12 = self.y12s[i_emb]
            y13 = self.y13s[i_emb]
            self.cos_sim_12 = (T.sum(layer_input1 * layer_input2, axis=-1) /
                               (layer_input1.norm(2, axis=-1) *
                                layer_input2.norm(2, axis=-1)))
            self.cos_sim_13 = (T.sum(layer_input1 * layer_input3, axis=-1) /
                               (layer_input1.norm(2, axis=-1) *
                                layer_input3.norm(2, axis=-1)))
            self.acos_sim_12 = T.arccos(self.cos_sim_12)
            self.acos_sim_13 = T.arccos(self.cos_sim_13)

            self.cos2_sim_12 = self.cos_sim_12**2
            self.cos2_sim_13 = self.cos_sim_13**2

            cos2_sim_cost = (y12+y13)*1. +\
                    T.switch(y12, self.cos2_sim_12, -self.cos2_sim_12) +\
                    T.switch(y13, self.cos2_sim_13, -self.cos2_sim_13)
            acos_sim_cost = (y12+y13)*1. +\
                    T.switch(y12, self.acos_sim_12, -self.acos_sim_12) +\
                    T.switch(y13, self.acos_sim_13, -self.acos_sim_13)
            self.cos2_costs.append(cos2_sim_cost)
            self.acos_costs.append(acos_sim_cost)

# TODO HERE
        self.cos2_sim_cost = T.sum(self.cos2_costs)
        self.mean_cos2_sim_cost = T.mean(self.cos2_sim_cost)
        self.sum_cos2_sim_cost = T.sum(self.cos2_sim_cost)

        self.acos_sim_cost = T.sum(self.acos_costs)
        self.mean_acos_sim_cost = T.mean(self.acos_sim_cost)
        self.sum_acos_sim_cost = T.sum(self.acos_sim_cost)

        if loss == 'cos2':
            self.cost = self.sum_cos_cos2_sim_cost
            self.mean_cost = self.mean_cos_cos2_sim_cost
        elif loss == 'acos':
            self.cost = self.sum_cos_sim_cost
            self.mean_cost = self.mean_cos_sim_cost
        else:
            print >> sys.stderr, "NO COST FUNCTION"
            sys.exit(-1)

        if debugprint:
            theano.printing.debugprint(self.cost)

        if hasattr(self, 'cost'):
            self.cost_training = self.cost
        if hasattr(self, 'mean_cost'):
            self.mean_cost_training = self.mean_cost
Esempio n. 34
0
a_ang_inner = tt.dscalar("a_ang_inner")  # milliarcsec

# the semi-major axis in au
a_inner = 1e-3 * a_ang_inner / parallax  # au

logP_inner = tt.dscalar("logP_inner")  # days
P_inner = tt.exp(logP_inner)

e_inner = tt.dscalar("e_inner")

omega_inner = tt.dscalar("omega_inner")  # omega_Aa
Omega_inner = tt.dscalar("Omega_inner")

cos_incl_inner = tt.dscalar("cos_incl_inner")
incl_inner = tt.arccos(cos_incl_inner)

MAb = tt.dscalar("MAb")

t_periastron_inner = tt.dscalar("t_periastron_inner")

orbit_inner = xo.orbits.KeplerianOrbit(
    a=a_inner * au_to_R_sun,
    period=P_inner,
    ecc=e_inner,
    t_periastron=t_periastron_inner,
    omega=omega_inner,
    Omega=Omega_inner,
    incl=incl_inner,
    m_planet=MAb,
)
Esempio n. 35
0
File: abcnet.py Progetto: RolT/abnet
    def __init__(
        self,
        numpy_rng,
        theano_rng=None,
        n_ins=40 * 3,
        layers_types=[ReLU, ReLU, ReLU, ReLU, ReLU],
        layers_sizes=[1024, 1024, 1024, 1024],
        n_outs=100,
        n_embs=2,
        loss="cos2",
        rho=0.9,
        eps=1.0e-6,
        max_norm=0.0,
        debugprint=False,
    ):
        self.layers = []
        self.params = []
        self.n_layers = len(layers_types)
        self.layers_types = layers_types
        assert self.n_layers > 0
        self.max_norm = max_norm
        self._rho = rho  # ``momentum'' for adadelta
        self._eps = eps  # epsilon for adadelta
        self._accugrads = []  # for adadelta
        self._accudeltas = []  # for adadelta

        if theano_rng == None:
            theano_rng = RandomStreams(numpy_rng.randint(2 ** 30))

        self.x1 = T.fmatrix("x1")
        self.x2 = T.fmatrix("x2")
        self.x3 = T.fmatrix("x3")
        self.y12s = [T.ivector("y12") for _ in n_embs]
        self.y13s = [T.ivector("y13") for _ in n_embs]
        self.n_embs = n_embs

        self.layers_ins = [n_ins] + layers_sizes[:-1]
        self.layers_outs = layers_sizes
        layer_input1 = self.x1
        layer_input2 = self.x2
        layer_input3 = self.x3
        layer_embs1 = [None for _ in n_embs]
        layer_embs2 = [None for _ in n_embs]
        layer_embs3 = [None for _ in n_embs]

        for layer_ind, (layer_type, n_in, n_out) in enumerate(
            zip(layers_types[:-1], self.layers_ins, self.layers_outs)
        ):
            this_layer1 = layer_type(rng=numpy_rng, input=layer_input1, n_in=n_in, n_out=n_out)  # , cap=6.)
            assert hasattr(this_layer1, "output")
            layer_input1 = this_layer1.output
            self.params.extend(this_layer1.params)
            self._accugrads.extend([build_shared_zeros(t.shape.eval(), "accugrad") for t in this_layer1.params])
            self._accudeltas.extend([build_shared_zeros(t.shape.eval(), "accudelta") for t in this_layer1.params])
            self.layers.append(this_layer1)
            this_layer2 = layer_type(
                rng=numpy_rng, input=layer_input2, n_in=n_in, n_out=n_out, W=this_layer1.W, b=this_layer1.b
            )  # , cap=6.)
            assert hasattr(this_layer2, "output")
            this_layer3 = layer_type(
                rng=numpy_rng, input=layer_input2, n_in=n_in, n_out=n_out, W=this_layer1.W, b=this_layer1.b
            )  # , cap=6.)
            assert hasattr(this_layer3, "output")
            layer_input2 = this_layer2.output
            self.layers.append(this_layer2)
            layer_input3 = this_layer3.output
            self.layers.append(this_layer3)

        for i_emb in xrange(n_embs):
            emb_layer_type = layers_types[-1]
            n_in = layers_sizes[-1]
            n_out = n_outs
            this_emb_layer1 = emb_layer_type(rng=numpy_rng, input=layer_input1, n_in=n_in, n_out=n_out)  # , cap=6.)
            assert hasattr(this_emb_layer1, "output")
            layer_embs1[i_emb] = this_emb_layer1.output
            self.params.extend(this_emb_layer1.params)
            self._accugrads.extend([build_shared_zeros(t.shape.eval(), "accugrad") for t in this_emb_layer1.params])
            self._accudeltas.extend([build_shared_zeros(t.shape.eval(), "accudelta") for t in this_emb_layer1.params])
            this_emb_layer2 = emb_layer_type(
                rng=numpy_rng, input=layer_input2, n_in=n_in, n_out=n_out, W=this_emb_layer1.W, b=this_emb_layer1.b
            )  # , cap=6.)
            assert hasattr(this_emb_layer2, "output")
            layer_embs2[i_emb] = this_emb_layer2.output
            this_emb_layer3 = emb_layer_type(
                rng=numpy_rng, input=layer_input3, n_in=n_in, n_out=n_out, W=this_emb_layer1.W, b=this_emb_layer1.b
            )  # , cap=6.)
            assert hasattr(this_emb_layer3, "output")
            layer_embs3[i_emb] = this_emb_layer3.output
            self.layers.append(this_emb_layer1)
            self.layers.append(this_emb_layer2)
            self.layers.append(this_emb_layer3)

        L2 = 0.0
        for param in self.params:
            L2 += T.sum(param ** 2)
        L1 = 0.0
        for param in self.params:
            L1 += T.sum(abs(param))

        self.cos2_costs = []
        self.acos_costs = []

        for i_emb in xrange(n_embs):
            y12 = self.y12s[i_emb]
            y13 = self.y13s[i_emb]
            self.cos_sim_12 = T.sum(layer_input1 * layer_input2, axis=-1) / (
                layer_input1.norm(2, axis=-1) * layer_input2.norm(2, axis=-1)
            )
            self.cos_sim_13 = T.sum(layer_input1 * layer_input3, axis=-1) / (
                layer_input1.norm(2, axis=-1) * layer_input3.norm(2, axis=-1)
            )
            self.acos_sim_12 = T.arccos(self.cos_sim_12)
            self.acos_sim_13 = T.arccos(self.cos_sim_13)

            self.cos2_sim_12 = self.cos_sim_12 ** 2
            self.cos2_sim_13 = self.cos_sim_13 ** 2

            cos2_sim_cost = (
                (y12 + y13) * 1.0
                + T.switch(y12, self.cos2_sim_12, -self.cos2_sim_12)
                + T.switch(y13, self.cos2_sim_13, -self.cos2_sim_13)
            )
            acos_sim_cost = (
                (y12 + y13) * 1.0
                + T.switch(y12, self.acos_sim_12, -self.acos_sim_12)
                + T.switch(y13, self.acos_sim_13, -self.acos_sim_13)
            )
            self.cos2_costs.append(cos2_sim_cost)
            self.acos_costs.append(acos_sim_cost)

        # TODO HERE
        self.cos2_sim_cost = T.sum(self.cos2_costs)
        self.mean_cos2_sim_cost = T.mean(self.cos2_sim_cost)
        self.sum_cos2_sim_cost = T.sum(self.cos2_sim_cost)

        self.acos_sim_cost = T.sum(self.acos_costs)
        self.mean_acos_sim_cost = T.mean(self.acos_sim_cost)
        self.sum_acos_sim_cost = T.sum(self.acos_sim_cost)

        if loss == "cos2":
            self.cost = self.sum_cos_cos2_sim_cost
            self.mean_cost = self.mean_cos_cos2_sim_cost
        elif loss == "acos":
            self.cost = self.sum_cos_sim_cost
            self.mean_cost = self.mean_cos_sim_cost
        else:
            print >>sys.stderr, "NO COST FUNCTION"
            sys.exit(-1)

        if debugprint:
            theano.printing.debugprint(self.cost)

        if hasattr(self, "cost"):
            self.cost_training = self.cost
        if hasattr(self, "mean_cost"):
            self.mean_cost_training = self.mean_cost
Esempio n. 36
0

def bound_loss(x, tnp=np):
    eps = 1e-9
    loss = tnp.maximum(tnp.maximum(eps, x - 1), tnp.maximum(eps, -x)) + eps
    return tnp.maximum(loss, eps) + eps


params_plus1 = shared(np.random.rand(nbatch, 2))
params_plus2 = shared(np.random.rand(nbatch, 2))

a = Print("invplus(x, params_plus1")(invplus(x, params_plus1))
b = Print("invplus(x, params_plus2")(invplus(y, params_plus2))

eps = 1e-9
a2 = T.arccos(T.clip(a, eps, 1 - eps))
b2 = T.arcsin(T.clip(b, eps, 1 - eps))
bl1 = bound_loss(a, tnp=T)
bl2 = bound_loss(b, tnp=T)

phi1_group = []
phi2_group = []
phi3_group = []
phi1_group.append(a2[:, 0])
phi1_group.append(b2[:, 0])

delta_group = []
delta_group.append(a2[:, 1])
delta_group.append(b2[:, 1])
delta_single, delta_var = singular(delta_group, name="delta")
Esempio n. 37
0
def _get_compiled_theano_functions(N_QUAD_PTS):
    # Planet masses: m1,m2
    m1, m2 = T.dscalars(2)
    mstar = 1
    mu1 = m1 * mstar / (mstar + m1)
    mu2 = m2 * mstar / (mstar + m2)
    Mstar1 = mstar + m1
    Mstar2 = mstar + m2
    beta1 = mu1 * T.sqrt(Mstar1 / mstar) / (mu1 + mu2)
    beta2 = mu2 * T.sqrt(Mstar2 / mstar) / (mu1 + mu2)
    j, k = T.lscalars('jk')
    s = (j - k) / k

    # Angle variable for averaging over
    psi = T.dvector()

    # Quadrature weights
    quad_weights = T.dvector('w')

    # Dynamical variables:
    Ndof = 3
    Nconst = 1
    dyvars = T.vector()
    y1, y2, y_inc, x1, x2, x_inc, amd = [
        dyvars[i] for i in range(2 * Ndof + Nconst)
    ]

    a20 = T.constant(1.)
    a10 = ((j - k) / j)**(2 / 3) * (Mstar1 / Mstar2)**(1 / 3)
    L10 = beta1 * T.sqrt(a10)
    L20 = beta2 * T.sqrt(a20)
    Ltot = L10 + L20
    f = L10 / L20
    L2res = (Ltot + amd) / (1 + f)
    Psi = -k * (s * L2res + (1 + s) * f * L2res)
    ###
    # actions
    ###
    I1 = 0.5 * (x1 * x1 + y1 * y1)
    I2 = 0.5 * (x2 * x2 + y2 * y2)
    Phi = 0.5 * (x_inc * x_inc + y_inc * y_inc)
    L1 = -s * Ltot - Psi / k - s * (I1 + I2 + Phi)
    L2 = (1 + s) * Ltot + Psi / k + (1 + s) * (I1 + I2 + Phi)

    # Set lambda2=0
    l2 = T.constant(0.)
    l1 = -1 * k * psi
    theta_res = (1 + s) * l2 - s * l1
    cos_theta_res = T.cos(theta_res)
    sin_theta_res = T.sin(theta_res)

    kappa1 = x1 * cos_theta_res + y1 * sin_theta_res
    eta1 = y1 * cos_theta_res - x1 * sin_theta_res

    kappa2 = x2 * cos_theta_res + y2 * sin_theta_res
    eta2 = y2 * cos_theta_res - x2 * sin_theta_res

    sigma = x_inc * cos_theta_res + y_inc * sin_theta_res
    rho = y_inc * cos_theta_res - x_inc * sin_theta_res
    # y = (sigma-i*rho)/sqrt(2)
    #   = sqrt(Phi) * exp[i (Omega1+Omega2) / 2]
    # Malige+ 2002,  Eqs 20 and 21
    r2byr1 = (L2 - L1 - I2 + I1) / Ltot
    sigma1 = rho * T.sqrt(1 + r2byr1) / T.sqrt(2)
    sigma2 = -rho * T.sqrt(1 - r2byr1) / T.sqrt(2)
    rho1 = -sigma * T.sqrt(1 + r2byr1) / T.sqrt(2)
    rho2 = sigma * T.sqrt(1 - r2byr1) / T.sqrt(2)

    Xre1 = kappa1 / T.sqrt(L1)
    Xim1 = -eta1 / T.sqrt(L1)
    Yre1 = 0.5 * sigma1 / T.sqrt(L1)
    Yim1 = -0.5 * rho1 / T.sqrt(L1)

    Xre2 = kappa2 / T.sqrt(L2)
    Xim2 = -eta2 / T.sqrt(L2)
    Yre2 = 0.5 * sigma2 / T.sqrt(L2)
    Yim2 = -0.5 * rho2 / T.sqrt(L2)

    absX1_sq = 2 * I1 / L1
    absX2_sq = 2 * I2 / L2
    X_to_z1 = T.sqrt(1 - absX1_sq / 4)
    X_to_z2 = T.sqrt(1 - absX2_sq / 4)
    Y_to_zeta1 = 1 / T.sqrt(1 - absX1_sq / 2)
    Y_to_zeta2 = 1 / T.sqrt(1 - absX2_sq / 2)

    a1 = (L1 / beta1)**2
    k1 = Xre1 * X_to_z1
    h1 = Xim1 * X_to_z1
    q1 = Yre1 * Y_to_zeta1
    p1 = Yim1 * Y_to_zeta1
    e1 = T.sqrt(absX1_sq) * X_to_z1
    inc1 = 2 * T.arcsin(T.sqrt(p1 * p1 + q1 * q1))

    a2 = (L2 / beta2)**2
    k2 = Xre2 * X_to_z2
    h2 = Xim2 * X_to_z2
    q2 = Yre2 * Y_to_zeta2
    p2 = Yim2 * Y_to_zeta2
    e2 = T.sqrt(absX2_sq) * X_to_z2
    inc2 = 2 * T.arcsin(T.sqrt(p2 * p2 + q2 * q2))

    beta1p = T.sqrt(Mstar1) * beta1
    beta2p = T.sqrt(Mstar2) * beta2
    Hkep = -0.5 * beta1p / a1 - 0.5 * beta2p / a2

    Hdir, Hind = calc_Hint_components_spatial(a1, a2, l1, l2, h1, k1, h2, k2,
                                              p1, q1, p2, q2, Mstar1, Mstar2)
    eps = m1 * m2 / (mu1 + mu2) / T.sqrt(mstar)
    Hpert = (Hdir + Hind / mstar)
    Hpert_av = Hpert.dot(quad_weights)
    Htot = Hkep + eps * Hpert_av

    #####################################################
    # Set parameters for compiling functions with Theano
    #####################################################

    # Get numerical quadrature nodes and weights
    nodes, weights = np.polynomial.legendre.leggauss(N_QUAD_PTS)

    # Rescale for integration interval from [-1,1] to [-pi,pi]
    nodes = nodes * np.pi
    weights = weights * 0.5

    # 'givens' will fix some parameters of Theano functions compiled below
    givens = [(psi, nodes), (quad_weights, weights)]

    # 'ins' will set the inputs of Theano functions compiled below
    #   Note: 'extra_ins' will be passed as values of object attributes
    #   of the 'ResonanceEquations' class 'defined below
    extra_ins = [m1, m2, j, k]
    ins = [dyvars] + extra_ins

    Stilde = Phi * (L2 - I2 - L1 + I1) / (Ltot)
    Q1 = 0.5 * (Phi + Stilde)
    Q2 = 0.5 * (Phi - Stilde)
    inc1 = T.arccos(1 - Q1 / (L1 - I1))
    inc2 = T.arccos(1 - Q2 / (L2 - I2))

    orbels = [
        a1, e1, inc1, k * T.arctan2(y1, x1), a2, e2, inc2,
        k * T.arctan2(y2, x2),
        T.arctan2(y_inc, x_inc)
    ]
    orbels_dict = dict(
        zip([
            'a1', 'e1', 'inc1', 'theta1', 'a2', 'e2', 'inc2', 'theta2', 'phi'
        ], orbels))

    actions = [L1, L2, I1, I2, Q1, Q2]
    actions_dict = dict(
        zip(['L1', 'L2', 'Gamma1', 'Gamma2', 'Q1', 'Q2'], actions))

    #  Conservative flow
    gradHtot = T.grad(Htot, wrt=dyvars)
    gradHpert = T.grad(Hpert_av, wrt=dyvars)
    gradHkep = T.grad(Hkep, wrt=dyvars)

    hessHtot = theano.gradient.hessian(Htot, wrt=dyvars)
    hessHpert = theano.gradient.hessian(Hpert_av, wrt=dyvars)
    hessHkep = theano.gradient.hessian(Hkep, wrt=dyvars)

    Jtens = T.as_tensor(np.pad(getOmegaMatrix(Ndof), (0, Nconst), 'constant'))
    H_flow_vec = Jtens.dot(gradHtot)
    Hpert_flow_vec = Jtens.dot(gradHpert)
    Hkep_flow_vec = Jtens.dot(gradHkep)

    H_flow_jac = Jtens.dot(hessHtot)
    Hpert_flow_jac = Jtens.dot(hessHpert)
    Hkep_flow_jac = Jtens.dot(hessHkep)

    ##########################
    # Compile Theano functions
    ##########################
    func_dict = {
        # Hamiltonians
        'H': Htot,
        #'Hpert':Hpert_av,
        #'Hkep':Hkep,
        ## Hamiltonian flows
        'H_flow': H_flow_vec,
        #'Hpert_flow':Hpert_flow_vec,
        #'Hkep_flow':Hkep_flow_vec,
        ## Hamiltonian flow Jacobians
        'H_flow_jac': H_flow_jac,
        #'Hpert_flow_jac':Hpert_flow_jac,
        #'Hkep_flow_jac':Hkep_flow_jac,
        ## Extras
        'orbital_elements': orbels_dict,
        'actions': actions_dict
    }
    compiled_func_dict = dict()
    with tqdm(func_dict.items()) as t:
        for key, val in t:
            t.set_description("Compiling '{}'".format(key))
            if key is 'timescales':
                inputs = extra_ins
            else:
                inputs = ins
            cf = theano.function(inputs=inputs,
                                 outputs=val,
                                 givens=givens,
                                 on_unused_input='ignore')
            compiled_func_dict[key] = cf
    return compiled_func_dict
Esempio n. 38
0
    def evaluate(self, dt, positions, velocities, rot_matrices, motor_signals):

        # ALL CONSTRAINTS CAN BE TRANSFORMED TO VELOCITY CONSTRAINTS!
        ##################
        # --- Step 1 --- #
        ##################
        # First, we integrate the applied force F_a acting of each rigid body (like gravity, ...) and
        # we obtain some new velocities v2' that tends to violate the constraints.

        totalforce = np.array([0,0,0,0,0,0], dtype='float32')  # total force acting on body outside of constraints
        acceleration = np.array([0,0,-9.81,0,0,0], dtype='float32')  # acceleration of the default frame
        newv = velocities + dt * acceleration[None,None,:]
        originalv = newv


        ##################
        # --- Step 2 --- #
        ##################
        # now enforce the constraints by having corrective impulses
        # convert mass matrices to world coordinates
        M = T.zeros(shape=(self.batch_size,self.num_bodies,6,6))


        M00 = self.lower_inertia_inv
        M01 = T.zeros(shape=(self.batch_size,self.num_bodies,3,3))
        M10 = M01
        M11 = T.sum(rot_matrices[:,:,:,None,:,None] * self.upper_inertia_inv[:,:,:,:,None,None] * rot_matrices[:,:,None,:,None,:], axis=(2,3))
        M0 = T.concatenate([M00,M01],axis=3)
        M1 = T.concatenate([M10,M11],axis=3)
        M = T.concatenate([M0,M1],axis=2)

        # changes every timestep

        # constraints are first dimension! We will need to stack them afterwards!
        J = [np.zeros((self.batch_size,6), dtype="float32") for _ in xrange(2 * self.num_constraints)]   # 0 constraints x 0 bodies x 2 objects x 6 states
        b_res =   [np.zeros((self.batch_size,), dtype="float32") for _ in xrange(self.num_constraints)]  # 0 constraints x 0 bodies
        b_error = [np.zeros((self.batch_size,), dtype="float32") for _ in xrange(self.num_constraints)]  # 0 constraints x 0 bodies
        C =       [np.zeros((self.batch_size,), dtype="float32") for _ in xrange(self.num_constraints)]  # 0 constraints x 0 bodies

        c_idx = 0
        for constraint,references,parameters in self.constraints:
            idx1 = references[0]
            if references[1] is not None:
                idx2 = references[1]
            else:
                idx2 = None

            if constraint == "universe":
                batched_J = numpy_repeat_new_axis(np.concatenate([-np.eye(3), np.zeros((3,3))]), self.batch_size).astype('float32')

                for i in xrange(3):
                    J[2*(c_idx+i)+0] = batched_J[:,:,i]
                    #J[c_idx+i,1,:] = np.concatenate([ np.eye(3), np.zeros((3,3))])[:,i]
                    b_error[c_idx+i] = -positions[:,idx1,i]
                c_idx += 3

                batched_J = numpy_repeat_new_axis(np.concatenate([np.zeros((3,3)),-np.eye(3)]), self.batch_size).astype('float32')
                for i in xrange(3):
                    J[2*(c_idx+i)+0] = batched_J[:,:,i]
                    #J[c_idx+i,1,:] = np.concatenate([np.zeros((3,3), dtype=DTYPE), np.eye(3)])[:,i]
                    b_error[c_idx+i] = np.zeros(shape=(self.batch_size,))
                c_idx += 3

            if constraint == "ball-and-socket" or constraint == "hinge" or constraint == "fixed":
                r1x = theano_convert_model_to_world_coordinate_no_bias(parameters["joint_in_model1_coordinates"][None,:], rot_matrices[:,idx1,:,:])
                r2x = theano_convert_model_to_world_coordinate_no_bias(parameters["joint_in_model2_coordinates"][None,:], rot_matrices[:,idx2,:,:])
                ss_r1x = single_skew_symmetric(r1x)
                ss_r2x = single_skew_symmetric(r2x)
                batched_eye = numpy_repeat_new_axis(np.eye(3, dtype='float32'), self.batch_size)
                complete_J1 = T.concatenate([-batched_eye, ss_r1x],axis=1)
                complete_J2 = T.concatenate([ batched_eye,-ss_r2x],axis=1)
                error = positions[:,idx2,:]+r2x-positions[:,idx1,:]-r1x
                for i in xrange(3):
                    J[2*(c_idx+i)+0] = complete_J1[:,:,i]
                    J[2*(c_idx+i)+1] = complete_J2[:,:,i]

                    b_error[c_idx+i] = error[:,i]
                c_idx += 3

            if constraint == "slider" or constraint == "fixed":
                batched_eye = numpy_repeat_new_axis(np.eye(3, dtype='float32'), self.batch_size)
                batched_zeros = numpy_repeat_new_axis(np.zeros((3,3), dtype='float32'), self.batch_size)


                complete_J1 = np.concatenate([batched_zeros,-batched_eye],axis=1)
                complete_J2 = np.concatenate([batched_zeros, batched_eye],axis=1)

                for i in xrange(3):
                    J[2*(c_idx+i)+0] = complete_J1[:,:,i]
                    J[2*(c_idx+i)+1] = complete_J2[:,:,i]

                rot_current = np.dot(rot_matrices[idx2,:,:], rot_matrices[idx1,:,:].T)
                rot_diff = np.dot(rot_current, parameters['rot_init'].T)
                cross = rot_diff.T - rot_diff
                # TODO: add stabilization of this constraint
                b_error[c_idx] = np.zeros(shape=(self.batch_size,))#cross[1,2]
                b_error[c_idx+1] = np.zeros(shape=(self.batch_size,))#cross[2,0]
                b_error[c_idx+2] = np.zeros(shape=(self.batch_size,))#cross[0,1]
                c_idx += 3

            if constraint == "hinge":
                a2x = theano_convert_model_to_world_coordinate_no_bias(parameters['axis_in_model2_coordinates'][None,:], rot_matrices[:,idx2,:,:])
                b1x = theano_convert_model_to_world_coordinate_no_bias(parameters['axis1_in_model1_coordinates'][None,:], rot_matrices[:,idx1,:,:])
                c1x = theano_convert_model_to_world_coordinate_no_bias(parameters['axis2_in_model1_coordinates'][None,:], rot_matrices[:,idx1,:,:])
                ss_a2x = single_skew_symmetric(a2x)

                batched_zeros = numpy_repeat_new_axis(np.zeros((3,), dtype='float32'), self.batch_size)

                J[2*(c_idx+0)+0] = T.concatenate([batched_zeros,-theano_dot_last_dimension_vector_matrix(b1x,ss_a2x)],axis=1)
                J[2*(c_idx+0)+1] = T.concatenate([batched_zeros, theano_dot_last_dimension_vector_matrix(b1x,ss_a2x)],axis=1)
                J[2*(c_idx+1)+0] = T.concatenate([batched_zeros,-theano_dot_last_dimension_vector_matrix(c1x,ss_a2x)],axis=1)
                J[2*(c_idx+1)+1] = T.concatenate([batched_zeros, theano_dot_last_dimension_vector_matrix(c1x,ss_a2x)],axis=1)

                b_error[c_idx+0] = theano_dot_last_dimension_vectors(a2x,b1x)
                b_error[c_idx+1] = theano_dot_last_dimension_vectors(a2x,c1x)
                c_idx += 2

            """
            if constraint == "limit":
                #TODO: move to Theano...
                angle = parameters["angle"]/180. * np.pi
                a = theano_convert_model_to_world_coordinate_no_bias(parameters['axis_in_model1_coordinates'], rot_matrices[:,idx1,:,:])
                rot_current = theano_dot_last_dimension_matrices(rot_matrices[:,idx2,:,:], rot_matrices[idx1,:,:].dimshuffle(0,2,1))
                rot_diff = theano_dot_last_dimension_matrices(rot_current, quat_to_rot_matrix(parameters['rot_init']).dimshuffle(0,2,1))
                theta2 = np.arccos(0.5*(np.trace(rot_diff)-1))
                cross = rot_diff.T - rot_diff
                dot2 = cross[1,2] * a[0] + cross[2,0] * a[1] + cross[0,1] * a[2]
                theta = ((dot2>0) * 2 - 1) * theta2

                if parameters["angle"] < 0:
                    J[2*c_idx+0] = np.concatenate([np.zeros((3,), dtype='float32'),-a])
                    J[2*c_idx+1] = np.concatenate([np.zeros((3,), dtype='float32'), a])
                else:
                    J[2*c_idx+0] = np.concatenate([np.zeros((3,), dtype='float32'), a])
                    J[2*c_idx+1] = np.concatenate([np.zeros((3,), dtype='float32'),-a])

                b_error[c_idx] = np.abs(angle - theta)
                if parameters["angle"] > 0:
                    b_error[c_idx] = angle - theta
                    C[c_idx] = (theta > angle)
                else:
                    b_error[c_idx] = theta - angle
                    C[c_idx] = (theta < angle)

                c_idx += 1
            """
            if constraint == "motor":
                ac = parameters['axis_in_model2_coordinates'][None,:]
                a = theano_convert_model_to_world_coordinate_no_bias(ac, rot_matrices[:,idx2,:,:])

                # TODO: remove dimshuffle(0,2,1) by using batched_dot
                rot_current = theano_dot_last_dimension_matrices(rot_matrices[:,idx2,:,:], rot_matrices[:,idx1,:,:].dimshuffle(0,2,1))
                rot_init = numpy_repeat_new_axis(parameters['rot_init'].T, self.batch_size)
                rot_diff = theano_dot_last_dimension_matrices(rot_current, rot_init)

                traces = rot_diff[:,0,0] + rot_diff[:,1,1] + rot_diff[:,2,2]

                # grad when x=-1 or x=1 does not exist for arccos
                theta2 = T.arccos(T.clip(0.5*(traces-1),-1+eps,1-eps))
                cross = rot_diff.dimshuffle(0,2,1) - rot_diff
                dot2 = cross[:,1,2] * ac[:,0] + cross[:,2,0] * ac[:,1] + cross[:,0,1] * ac[:,2]

                theta = ((dot2>0) * 2 - 1) * theta2

                batched_zeros = numpy_repeat_new_axis(np.zeros((3,), dtype='float32'), self.batch_size)
                J[2*c_idx+0] = T.concatenate([batched_zeros,-a],axis=1)
                J[2*c_idx+1] = T.concatenate([batched_zeros, a],axis=1)

                motor_signal = motor_signals[:,parameters["motor_id"]]

                if "min" in parameters and "max" in parameters:
                    motor_min = (parameters["min"]/180. * np.pi)
                    motor_max = (parameters["max"]/180. * np.pi)
                    motor_signal = T.clip(motor_signal, motor_min, motor_max).astype('float32')

                def smallestSignedAngleBetween(x, y):
                    a1 = (x - y) % np.float32(2*np.pi)
                    b1 = (y - x) % np.float32(2*np.pi)
                    return T.minimum(a1,b1)*((a1>b1)*2-1)

                error_signal = -smallestSignedAngleBetween(theta, motor_signal)

                if parameters["type"] == "velocity":
                    b_error[c_idx] = motor_signal
                elif parameters["type"] == "position":
                    if "delta" in parameters  and "motor_velocity" in parameters:
                        velocity = parameters["motor_velocity"] / 180. * np.pi
                        b_error[c_idx] = dt * T.clip((abs(error_signal) > parameters["delta"]) * error_signal * parameters["motor_gain"], -velocity, velocity)
                    else:
                        b_error[c_idx] = dt * error_signal * parameters["motor_gain"]

                c_idx += 1

            if constraint == "ground":
                r = self.radii[idx1].astype('float32')
                J[2*c_idx+0] = numpy_repeat_new_axis(np.array([0,0,1,0,0,0], dtype='float32'), self.batch_size)
                J[2*c_idx+1] = numpy_repeat_new_axis(np.array([0,0,0,0,0,0], dtype='float32'), self.batch_size)

                b_error[c_idx] = T.clip(positions[:,idx1,Z] - r + parameters["delta"], np.finfo('float32').min, 0)
                b_res[c_idx] = parameters["alpha"] * newv[:,idx1,Z]
                C[c_idx] = positions[:,idx1,Z] - r
                c_idx += 1

            if constraint == "ground" and parameters["mu"]!=0:
                r = self.radii[idx1].astype('float32')
                for i in xrange(2):
                    if i==0:
                        J[2*(c_idx+i)+0] = numpy_repeat_new_axis(np.array([0,1,0,-r,0,0], dtype='float32'), self.batch_size)
                    else:
                        J[2*(c_idx+i)+0] = numpy_repeat_new_axis(np.array([1,0,0,0,r,0], dtype='float32'), self.batch_size)
                    J[2*(c_idx+i)+1] = numpy_repeat_new_axis(np.array([0,0,0,0,0,0], dtype='float32'), self.batch_size)
                    C[c_idx+i] = positions[:,idx1,Z] - r
                c_idx += 2

            if constraint == "ground" and parameters["torsional_friction"] and parameters["mu"]!=0:
                r = self.radii[idx1].astype('float32')
                J[2*c_idx+0] = numpy_repeat_new_axis(np.array([0,0,0,0,0,r], dtype='float32'), self.batch_size)
                J[2*c_idx+1] = numpy_repeat_new_axis(np.array([0,0,0,0,0,0], dtype='float32'), self.batch_size)
                C[c_idx] = positions[:,idx1,Z] - r
                c_idx += 1


        zipped_indices = [j for i in zip(self.zero_index,self.one_index) for j in i]
        mass_matrix = M[:,zipped_indices,:,:].reshape(shape=(self.batch_size, self.num_constraints, 2, 6, 6))

        #mass_matrix = T.concatenate((M[:,self.zero_index,None,:,:], M[:,self.one_index,None,:,:]), axis=2)
        #mass_matrix = T.concatenate((
        #                            T.stack([M[:,i,None,:,:] for i in self.zero_index], axis=1),
        #                            T.stack([M[:,j,None,:,:] for j in self.one_index], axis=1)
        #                            ), axis=2)

        #J = T.stack(J, axis=0).reshape(shape=(self.num_constraints,2,self.batch_size,6)).dimshuffle(2,0,1,3)

        J = theano_stack_batched_integers_mixed_numpy(J, expected_shape=(self.num_constraints,2,self.batch_size,6)).dimshuffle(2,0,1,3)

        C = theano_stack_batched_integers_mixed_numpy(C, expected_shape=(self.num_constraints,self.batch_size)).dimshuffle(1,0)
        b_res = theano_stack_batched_integers_mixed_numpy(b_res, expected_shape=(self.num_constraints,self.batch_size)).dimshuffle(1,0)
        b_error = theano_stack_batched_integers_mixed_numpy(b_error, expected_shape=(self.num_constraints,self.batch_size)).dimshuffle(1,0)

        #v = np.zeros((self.num_constraints,2,6))  # 0 constraints x 2 objects x 6 states

        self.P = self.warm_start * self.P

        for iteration in xrange(self.projected_gauss_seidel_iterations):
            # changes every iteration
            v = newv[:,zipped_indices,:].reshape(shape=(self.batch_size, self.num_constraints, 2, 6))
            #v = T.concatenate((newv[self.zero_index,None,:],newv[self.one_index,None,:]),axis=1)
            #v = T.concatenate((
            #    T.stack([newv[:,i,None,:] for i in self.zero_index], axis=1),
            #    T.stack([newv[:,j,None,:] for j in self.one_index], axis=1)
            #),axis=2)

            # TODO: batch-dot-product
            m_eff = 1./T.sum(T.sum(J[:,:,:,None,:]*mass_matrix, axis=4)*J, axis=(2,3))
            #m_eff = 1./T.sum(T.sum(J[:,:,None,:]*mass_matrix, axis=-1)*J, axis=(-1,-2))

            k = m_eff * (self.w**2)
            c = m_eff * 2*self.zeta*self.w

            CFM = 1./(c+dt*k)
            ERP = dt*k/(c+dt*k)

            m_c = 1./(1./m_eff + CFM)

            b = ERP/dt * b_error + b_res

            lamb = - m_c * (T.sum(J*v, axis=(2,3)) + CFM * self.P + b)

            self.P += lamb
            #print J[[39,61,65,69],:]
            #print np.sum(lamb**2), np.sum(self.P**2)

            #clipping_force = T.concatenate([self.P[:,j].dimshuffle(0,'x') for j in self.clipping_idx],axis=1)
            clipping_force = self.P[:,self.clipping_idx]

            clipping_limit = abs(self.clipping_a * clipping_force + self.clipping_b * dt)
            self.P = T.clip(self.P,-clipping_limit, clipping_limit)
            applicable = (1.0*(C<=0)) * (1-(self.only_when_positive*(self.P<=0)))

            # TODO: batch-dot-product
            result = T.sum(mass_matrix*J[:,:,:,None,:], axis=4) * (self.P * applicable)[:,:,None,None]
            #result = theano_dot_last_dimension_matrix_vector(mass_matrix,J) * (self.P * applicable)[:,:,None,None]

            result = result.reshape((self.batch_size, 2*self.num_constraints, 6))

            r = []
            for i in xrange(len(self.map_object_to_constraint)):
                delta_v = T.sum(result[:,self.map_object_to_constraint[i],:], axis=1)
                r.append(delta_v)
            newv = newv + T.stack(r, axis=1)
        #print
        return newv
Esempio n. 39
0
    def __init__(self,
                 period=None,
                 a=None,
                 t0=None,
                 t_periastron=None,
                 incl=None,
                 b=None,
                 duration=None,
                 ecc=None,
                 omega=None,
                 Omega=None,
                 m_planet=0.0,
                 m_star=None,
                 r_star=None,
                 rho_star=None,
                 m_planet_units=None,
                 rho_star_units=None,
                 model=None,
                 contact_points_kwargs=None,
                 **kwargs):
        add_citations_to_model(self.__citations__, model=model)

        self.gcc_to_sun = ((constants.M_sun / constants.R_sun**3).to(
            u.g / u.cm**3).value)
        self.au_to_R_sun = (constants.au / constants.R_sun).value
        self.G_grav = constants.G.to(u.R_sun**3 / u.M_sun / u.day**2).value

        self.kepler_op = KeplerOp(**kwargs)

        # Parameters
        # self.period = tt.as_tensor_variable(period)
        self.m_planet = tt.as_tensor_variable(m_planet)
        if m_planet_units is not None:
            self.m_planet *= (1 * m_planet_units).to(u.M_sun).value

        self.a, self.period, self.rho_star, self.r_star, self.m_star = \
            self._get_consistent_inputs(a, period, rho_star, r_star, m_star,
                                        rho_star_units)
        self.m_total = self.m_star + self.m_planet

        self.n = 2 * np.pi / self.period
        self.a_star = self.a * self.m_planet / self.m_total
        self.a_planet = -self.a * self.m_star / self.m_total

        self.K0 = self.n * self.a / self.m_total

        # Set up the contact points calculation
        if contact_points_kwargs is None:
            contact_points_kwargs = dict()

        if Omega is None:
            self.Omega = None
        else:
            self.Omega = tt.as_tensor_variable(Omega)
            self.cos_Omega = tt.cos(self.Omega)
            self.sin_Omega = tt.sin(self.Omega)

        # Eccentricity
        self.contact_points_op = ContactPointsOp(**contact_points_kwargs)
        if ecc is None:
            self.ecc = None
            self.M0 = 0.5 * np.pi + tt.zeros_like(self.n)
            incl_factor = 1
        else:
            self.ecc = tt.as_tensor_variable(ecc)
            if omega is None:
                raise ValueError("both e and omega must be provided")
            self.omega = tt.as_tensor_variable(omega)

            self.cos_omega = tt.cos(self.omega)
            self.sin_omega = tt.sin(self.omega)

            opsw = 1 + self.sin_omega
            E0 = 2 * tt.arctan2(
                tt.sqrt(1 - self.ecc) * self.cos_omega,
                tt.sqrt(1 + self.ecc) * opsw)
            self.M0 = E0 - self.ecc * tt.sin(E0)

            ome2 = 1 - self.ecc**2
            self.K0 /= tt.sqrt(ome2)
            incl_factor = (1 + self.ecc * self.sin_omega) / ome2

        if b is not None:
            if incl is not None or duration is not None:
                raise ValueError("only one of 'incl', 'b', and 'duration' can "
                                 "be given")
            self.b = tt.as_tensor_variable(b)
            self.cos_incl = incl_factor * self.b * self.r_star / self.a_planet
            self.incl = tt.arccos(self.cos_incl)
        elif incl is not None:
            if duration is not None:
                raise ValueError("only one of 'incl', 'b', and 'duration' can "
                                 "be given")
            self.incl = tt.as_tensor_variable(incl)
            self.cos_incl = tt.cos(self.incl)
            self.b = self.a_planet * self.cos_incl / (incl_factor *
                                                      self.r_star)
        elif duration is not None:
            if self.ecc is None:
                raise ValueError("fitting with duration only works for "
                                 "eccentric orbits")
            self.duration = tt.as_tensor_variable(duration)
            c = tt.sin(np.pi * self.duration * incl_factor / self.period)
            c2 = c * c
            aor = self.a_planet / self.r_star
            esinw = self.ecc * self.sin_omega
            self.b = tt.sqrt(
                (aor**2 * c2 - 1) / (c2 * esinw**2 + 2 * c2 * esinw + c2 -
                                     self.ecc**4 + 2 * self.ecc**2 - 1))
            self.b *= (1 - self.ecc**2)
            self.cos_incl = incl_factor * self.b * self.r_star / self.a_planet
            self.incl = tt.arccos(self.cos_incl)
        else:
            zla = tt.zeros_like(self.a)
            self.incl = 0.5 * np.pi + zla
            self.cos_incl = zla
            self.b = zla

        if t0 is not None and t_periastron is not None:
            raise ValueError("you can't define both t0 and t_periastron")
        if t0 is None and t_periastron is None:
            t0 = 0.0

        if t0 is None:
            self.t_periastron = tt.as_tensor_variable(t_periastron)
            self.t0 = self.t_periastron + self.M0 / self.n
        else:
            self.t0 = tt.as_tensor_variable(t0)
            self.t_periastron = self.t0 - self.M0 / self.n

        self.tref = self.t_periastron

        self.sin_incl = tt.sin(self.incl)
Esempio n. 40
0
    # generate Omega_disk from range of samples
    Omega_disk = pm.Uniform(
        "OmegaDisk",
        lower=np.min(Omega_samples),
        upper=np.max(Omega_samples),
        testval=np.mean(Omega_samples),
    )

    disk_observed = tt.as_tensor_variable([MA, i_disk, Omega_disk])

    pm.MvNormal("obs_disk", mu=disk_mu, cov=disk_cov, observed=disk_observed)

    # calculate the mutual inclination as well
    theta = pm.Deterministic(
        "thetaADisk",
        tt.arccos(
            tt.cos(i_disk) * tt.cos(incl) +
            tt.sin(i_disk) * tt.sin(incl) * tt.cos(Omega_disk - Omega)))

# iterate through the list of free_RVs in the model to get things like
# ['logKAa_interval__', etc...] then use a regex to strip away
# the transformations (in this case, _interval__ and _angle__)
# \S corresponds to any character that is not whitespace
# https://docs.python.org/3/library/re.html
sample_vars = [re.sub("_\S*__", "", var.name) for var in model.free_RVs]

all_vars = [
    var.name for var in model.unobserved_RVs
    if ("_interval__" not in var.name) and ("_angle__" not in var.name) and (
        "_lowerbound__" not in var.name)
]
Esempio n. 41
0
    a1 = pm.Deterministic('a1',
                          sigmaA * Da1 + mod.A1(f1_, [numax, w, A, V1, V2]))
    a2 = pm.Deterministic('a2',
                          sigmaA * Da2 + mod.A2(f2_, [numax, w, A, V1, V2]))

    h0 = pm.Deterministic('h0', 2 * tt.sqr(a0) / np.pi / g0)
    h1 = pm.Deterministic('h1', 2 * tt.sqr(a1) / np.pi / g1)
    h2 = pm.Deterministic('h2', 2 * tt.sqr(a2) / np.pi / g2)

    # Mode splitting & model
    xsplit = pm.HalfNormal('xsplit',
                           sigma=2.0,
                           testval=init_m[9] * np.sin(init_m[10]))
    cosi = pm.Uniform('cosi', 0., 1., testval=np.cos(init_m[10]))

    i = pm.Deterministic('i', tt.arccos(cosi))
    split = pm.Deterministic('split', xsplit / tt.sin(i))

    b = BNormal('b', mu=1., sigma=.1, testval=1.)

    fit = mod.model([f0, f1, f2, g0, g1, g2, h0, h1, h2, split, i, b])

    like = pm.Gamma('like', alpha=1., beta=1. / fit, observed=p)

# In[20]:

with pm_model:
    trace = pm.sample(chains=4, target_accept=.99)

# In[21]:
Esempio n. 42
0
 def pred(self, output):
     W1 = 0.5
     W2 = 0.5
     pred = T.argmin(W1 * T.arccos(T.dot(output, mappings_vec.T)/(absolute(output)*absolute(mappings_vec, axis=1))) +
                     W2 * (absolute(mappings_vec, axis=1) - absolute(output)))
     return pred