Пример #1
0
def Esqr(cosmo, a):
    r"""Square of the scale factor dependent factor E(a) in the Hubble
  parameter.

  Parameters
  ----------
  a : array_like
      Scale factor

  Returns
  -------
  E^2 : ndarray, or float if input scalar
      Square of the scaling of the Hubble constant as a function of
      scale factor

  Notes
  -----

  The Hubble parameter at scale factor `a` is given by
  :math:`H^2(a) = E^2(a) H_o^2` where :math:`E^2` is obtained through
  Friedman's Equation (see :cite:`2005:Percival`) :

  .. math::

      E^2(a) = \Omega_m a^{-3} + \Omega_k a^{-2} + \Omega_{de} a^{f(a)}

  where :math:`f(a)` is the Dark Energy evolution parameter computed
  by :py:meth:`.f_de`.
  """
    return cosmo.Omega_m*np.power(a, -3) + cosmo.Omega_k*np.power(a, -2) + \
        cosmo.Omega_de*np.power(a, f_de(cosmo, a))
Пример #2
0
def ft_phase_screen(r0, N, delta, L0, l0):
    # Set PSD
    del_f = 1/(N*delta)
    fx = []
    tmp = []
    for i in range(512):
        tmp.append(-256 + i)
    for i in range(512):
        fx.append(tmp)
    fx = jnp.array(fx)
    # frequency grid [1/m]
    fy = jnp.transpose(fx)
    f, th = cart2pol(fx, fy)
    fm = 5.92/l0/(2*jnp.pi)
    f0 = 1/L0
    # modified von Karman atmospheric phase PSD
    PSD_phi = 0.023*jnp.power(r0, -5/3) * jnp.exp(-jnp.power((f/fm), 2))/ jnp.power((jnp.power(f, 2) + jnp.power(f0, 2)), 11/6)
    PSD_phi = np.array(PSD_phi)
    PSD_phi[256][256] = 0
    # random draws of Fourier coefficients
    A = np.random.randn(N, N)
    B = np.random.randn(N, N)

    cn = (A + complex("j")*B) * jnp.sqrt(PSD_phi)*del_f
    phz = jnp.real(ift2(cn,1))
    return phz
Пример #3
0
    def _case1(zagf):
        z, alpha, _, flag = zagf

        # dz = - dCDF(z; a) / pdf(z; a)
        # pdf = z^(a-1) * e^(-z) / Gamma(a)
        # CDF(z; a) = IncompleteGamma(a, z) / Gamma(a)
        # dCDF(z; a) = (dIncompleteGamma - IncompleteGamma * Digamma(a)) / Gamma(a)
        #            =: unnormalized_dCDF / Gamma(a)
        # IncompleteGamma ~ z^a [ 1/a - z/(a+1) + z^2/2!(a+2) - z^3/3!(a+3) + z^4/4!(a+4) - z^5/5!(a+5) ]
        #                 =: z^a * term1
        # dIncompleteGamma ~ z^a * log(z) * term1 - z^a [1/a^2 - z/(a+1)^2 + z^2/2!(a+2)^2
        #                                                - z^3/3!(a+3)^2 + z^4/4!(a+4)^2 - z^5/5!(a+5)^2 ]
        #                  =: z^a * log(z) * term1 - z^a * term2
        # unnormalized_dCDF = z^a { [log(z) - Digamma(a)] * term1 - term2 }
        zi = 1.0
        update = zi / alpha
        term1 = update
        term2 = update / alpha
        for i in range(1, 6):
            zi = -zi * z / i
            update = zi / (alpha + i)
            term1 = term1 + update
            term2 = term2 + update / (alpha + i)

        unnormalized_cdf_dot = np.power(z, alpha) * (
            (np.log(z) - lax.digamma(alpha)) * term1 - term2)
        unnormalized_pdf = np.power(z, alpha - 1) * np.exp(-z)
        grad = -unnormalized_cdf_dot / unnormalized_pdf

        return z, alpha, grad, ~flag
Пример #4
0
 def objective(params: list, bparam: list, batch_input) -> float:
     result = 25.0
     for w1 in params:
         result += np.mean(
             np.divide(np.power(w1, 4), 4.0) +
             bparam[0] * np.divide(np.power(w1, 2), 2.0))
     return result
 def mu_init(key, shape, rng):
     low = -1 * 1 / jnp.power(x.shape[-1], 0.5)
     high = 1 * 1 / jnp.power(x.shape[-1], 0.5)
     return jax.random.uniform(rng,
                               shape=shape,
                               minval=low,
                               maxval=high)
Пример #6
0
    def getlambda(self, hyperparams=None):
        """Computing the diagonal of the Lambda matrix (spectral values)
        We're using jax.numpy routines since this function is used in the optimisation

        Parameters
        ----------
        hyperparams : ndarray
            array with hyperparameters, if not passed the object attributes are used

        Returns
        -------
        lambda : ndarray
            (m,) array with spectral values
        """

        # angular frequencies
        omega = self.mperms * 0.5 * jnp.pi / jnp.array([[self.Lx, self.Ly]])

        if hyperparams is None:
            l = self.l
            sf = self.sigma_f
        else:  # when optimising, we must be able to provide the parameters directly
            l = hyperparams[1]
            sf = hyperparams[0]

        # compute the spectral density values
        if self.covfunc.covtype == 'matern':  # matern
            return sf * jnp.power(
                2 * self.covfunc.nu / l + jnp.sum(omega**2, 1),
                -(self.covfunc.nu + 1)) / jnp.power(l, self.covfunc.nu)
        if self.covfunc.covtype == 'se':  # squared exponential
            return sf * l * jnp.exp(-0.5 * l + jnp.sum(omega**2, 1))
Пример #7
0
def pacs_model(priors):
    pointing_matrices = [([p.amat_row, p.amat_col], p.amat_data)
                         for p in priors]
    flux_lower = np.asarray([p.prior_flux_lower for p in priors]).T
    flux_upper = np.asarray([p.prior_flux_upper for p in priors]).T

    bkg_mu = np.asarray([p.bkg[0] for p in priors]).T
    bkg_sig = np.asarray([p.bkg[1] for p in priors]).T

    with numpyro.plate('bands', len(priors)):
        sigma_conf = numpyro.sample('sigma_conf', dist.HalfCauchy(1.0, 0.5))
        bkg = numpyro.sample('bkg', dist.Normal(bkg_mu, bkg_sig))

        with numpyro.plate('nsrc', priors[0].nsrc):
            src_f = numpyro.sample('src_f',
                                   dist.Uniform(flux_lower, flux_upper))
    db_hat_psw = sp_matmul(pointing_matrices[0], src_f[:, 0][:, None],
                           priors[0].snpix).reshape(-1) + bkg[0]
    db_hat_pmw = sp_matmul(pointing_matrices[1], src_f[:, 1][:, None],
                           priors[1].snpix).reshape(-1) + bkg[1]

    sigma_tot_psw = jnp.sqrt(
        jnp.power(priors[0].snim, 2) + jnp.power(sigma_conf[0], 2))
    sigma_tot_pmw = jnp.sqrt(
        jnp.power(priors[1].snim, 2) + jnp.power(sigma_conf[1], 2))

    with numpyro.plate('psw_pixels', priors[0].sim.size):  # as ind_psw:
        numpyro.sample("obs_psw",
                       dist.Normal(db_hat_psw, sigma_tot_psw),
                       obs=priors[0].sim)
    with numpyro.plate('pmw_pixels', priors[1].sim.size):  # as ind_pmw:
        numpyro.sample("obs_pmw",
                       dist.Normal(db_hat_pmw, sigma_tot_pmw),
                       obs=priors[1].sim)
Пример #8
0
    def T_seperate(params_all, x, dx, coeff_mask, z0):
        params_psi = params_all[:hyper_params['n_psi']]
        sindy_coeff = params_all[-1][0]

        z_opt = phi_vec(params_psi, x, z0)
        Theta = sindy_library(z_opt)

        dz = dz_func_vec(params_psi, x, dx, z0)

        dz_pred = dz_pred_vec(Theta, coeff_mask, sindy_coeff)
        x_rec = psi_vec(params_psi, z_opt)
        dx_rec = dx_network_vec(params_psi, z_opt, dz_pred)

        results = {}
        results['x_loss'] = jnp.mean(jnp.power(x - x_rec, 2))
        results['dx_loss'] = np.mean(jnp.power(dx - dx_rec, 2))
        results['dz_loss'] = jnp.mean(jnp.power(dz - dz_pred, 2))
        results['regul'] = jnp.mean(jnp.abs(sindy_coeff))

        dx_loss = jnp.multiply(hyper_params['eta1'], results['dx_loss'])
        dz_loss = jnp.multiply(hyper_params['eta2'], results['dz_loss'])
        regul = jnp.multiply(hyper_params['eta3'], results['regul'])

        results['loss'] = results['x_loss'] + dx_loss + dz_loss + regul
        results['x_rec'] = x_rec
        results['dx_rec'] = dx_rec
        results['z'] = z_opt
        results['dz'] = dz
        results['dz_pred'] = dz_pred
        return results
Пример #9
0
def quat2mat_jax(quat):
    norm_quat = quat
    norm_quat = norm_quat / jnp.linalg.norm(norm_quat, axis=1, keepdims=True)
    w, x, y, z = norm_quat[:, 0], norm_quat[:, 1], norm_quat[:,
                                                             2], norm_quat[:,
                                                                           3]

    batch_size = quat.shape[0]

    w2, x2, y2, z2 = jnp.power(w,
                               2), jnp.power(x,
                                             2), jnp.power(y,
                                                           2), jnp.power(z, 2)
    wx, wy, wz = w * x, w * y, w * z
    xy, xz, yz = x * y, x * z, y * z

    rotMat = jnp.stack(
        [
            w2 + x2 - y2 - z2,
            2 * xy - 2 * wz,
            2 * wy + 2 * xz,
            2 * wz + 2 * xy,
            w2 - x2 + y2 - z2,
            2 * yz - 2 * wx,
            2 * xz - 2 * wy,
            2 * wx + 2 * yz,
            w2 - x2 - y2 + z2,
        ],
        axis=1,
    ).reshape(batch_size, 3, 3)
    return rotMat
def loss_function(params, t, W, Xzero):  # idea take M,D out by making X0
    # X, Y, Y_tilde, Z, DYDT = vXYZpaths(params, t, W, Xzero)
    X, Y, Y_tilde, Z = vXYZpaths(params, t, W, Xzero)
    loss = jnp.sum(jnp.power(Y[:, 1:-1, :] - Y_tilde[:, 1:-1, :], 2))
    loss += jnp.sum(jnp.power(Y[:, -1, :] - vg_tf(X[:, -1, :]), 2))
    loss += jnp.sum(jnp.power(Z[:, -1, :] - vDg_tf(X[:, -1, :]), 2))
    return loss
Пример #11
0
def apply_regularization_loss(reg_type,
                              regularization_values,
                              regularization_norm=None):
    """Calculates regularization of specified type on supplied values.

  Args:
    reg_type: which regularization type found in regualarization_fn_dict to use
    regularization_values: values subject to regularization
    regularization_norm: <p> value if lp-norm is used

  Returns:
    regularization loss value

  Raises:
    ValueError: if the specified reg_type is not in regualarization_fn_dict
  """
    regualarization_fn_dict = {
        'l1_norm':
        lambda vals, _: sum(jnp.abs(vals)),
        'l2_norm':
        lambda vals, _: jnp.sqrt(sum(jnp.power(vals, 2))),
        'lp_norm':
        lambda vals, norm: jnp.power(  # pylint: disable=g-long-lambda
            sum(jnp.power(jnp.abs(vals), norm)), 1 / norm),
        '1_minus_mean_sqr':
        lambda vals, _: (1 - jnp.mean(vals))**2,
        '1_minus_mean':
        lambda vals, _: (1 - jnp.mean(vals))
    }
    if reg_type not in regualarization_fn_dict:
        raise ValueError(f'Unrecognized regularization type: {reg_type}.')
    reg_fn = regualarization_fn_dict[reg_type]
    return reg_fn(regularization_values, regularization_norm)
Пример #12
0
Файл: d3.py Проект: bast/d3
def d3(charges, coordinates, rs6, rs8, s6, s8, rab, rcov, r2r4, coefficients):
    bohr_to_angstrom = constants.conversion_factor("bohr", "angstrom")

    coordinates_angstrom = list(
        map(lambda x: x * bohr_to_angstrom, coordinates))
    coordination_numbers = compute_coordination_numbers(
        coordinates_angstrom, charges, rcov)

    result = 0.0
    for j, charge1 in enumerate(charges):
        for k, charge2 in enumerate(charges):
            if k > j:
                dx = coordinates[3 * j + 0] - coordinates[3 * k + 0]
                dy = coordinates[3 * j + 1] - coordinates[3 * k + 1]
                dz = coordinates[3 * j + 2] - coordinates[3 * k + 2]
                dist = jnp.sqrt(dx * dx + dy * dy + dz * dz)

                c6jk = get_c6jk(
                    coefficients[(charge1, charge2)],
                    coordination_numbers[j],
                    coordination_numbers[k],
                )

                rr = rab[(charge1, charge2)] / (dist * bohr_to_angstrom)
                r1 = (-s6 * c6jk / (jnp.power(dist, 6) *
                                    (1.0 + 6.0 * jnp.power(rs6 * rr, 14))))

                c8jk = 3.0 * c6jk * r2r4[charge1] * r2r4[charge2]
                r2 = (-s8 * c8jk / (jnp.power(dist, 8) *
                                    (1.0 + 6.0 * jnp.power(rs8 * rr, 16))))
                result += r1 + r2

    return result
Пример #13
0
 def mu_init(key, shape):
     # Initialization of mean noise parameters (Section 3.2)
     low = -1 / jnp.power(x.shape[0], 0.5)
     high = 1 / jnp.power(x.shape[0], 0.5)
     return jax.random.uniform(key,
                               minval=low,
                               maxval=high,
                               shape=shape)
Пример #14
0
 def _diag_grad_xy(
         self, x: jnp.ndarray, y: jnp.ndarray, c: Union[float, jnp.ndarray],
         beta: Union[float, jnp.ndarray],
         bandwidth: Union[float, jnp.ndarray]) -> Union[float, jnp.ndarray]:
     diff = (x - y) / bandwidth
     base = c + 0.5 * jnp.sum(jnp.square(diff), axis=-1)
     return (- beta * jnp.power(base, beta - 1) + beta * (beta - 1) * diff ** 2 * jnp.power(base, beta - 2)) \
            / bandwidth ** 2
def loss_function(params, t, W, Xzero):
    alpha = 1
    beta = 1
    X, Y, Y_tilde, Z = vXYZpaths(params, t, W, Xzero)
    loss = jnp.sum(jnp.power(Y[:, 1:-1, :] - Y_tilde[:, 1:-1, :], 2))
    loss += alpha * jnp.sum(jnp.power(Y[:, -1, :] - vg_tf(X[:, -1, :]), 2))
    loss += beta * jnp.sum(jnp.power(Z[:, -1, :] - vDg_tf(X[:, -1, :]), 2))
    return loss
Пример #16
0
def loss_function(params, t, W, Xzero):   #idea take M,D out by making X0
    X,Y,Y_tilde,Z = vXYZpaths(params, t, W, Xzero)
    # loss += jnp.sum(jnp.power(Y[:, :, 1:] - Y_tilde, 2)) # Y is 51, Y_tilde is 50 but only sum up to penultimate
    loss = jnp.sum(jnp.power(Y[:, 1:-1, :] - Y_tilde[:,  1:-1, :], 2)) + jnp.sum(jnp.power(Y[:,-1,:] - vg_tf(X[:,-1,:]), 2))
    loss += jnp.sum(jnp.power(Z[:,-1,:] - vDg_tf(X[:,-1,:]), 2)) #terminal 1st order condition remove for now
    # loss += jnp.sum(jnp.power(jnp.dot(Z[:,-1,:].T, sigma_tf(T, X[:,-1,:], Y[:,-1,:])) - vDg_tf(X[:,-1,:]), 2)) #terminal 1st order condition remove for now

    return loss
Пример #17
0
 def _tempred_cross_entropy_loss(unused_activations):
     loss_values = jnp.multiply(
         labels,
         log_t(labels + 1e-10, t1) -
         log_t(probabilities, t1)) - 1.0 / (2.0 - t1) * (jnp.power(
             labels, 2.0 - t1) - jnp.power(probabilities, 2.0 - t1))
     loss_values = jnp.sum(loss_values, -1)
     return loss_values
def loss_function(params, t, W, Xzero):  #idea take M,D out by making X0
    X, Y, Y_tilde = vXYpaths(params, t, W, Xzero)
    # loss += jnp.sum(jnp.power(Y[:, :, 1:] - Y_tilde, 2)) # Y is 51, Y_tilde is 50 but only sum up to penultimate
    # + terminal condition
    loss = jnp.sum(jnp.power(Y[:, :, 1:-2] - Y_tilde[:, :, :-2], 2)) + jnp.sum(
        jnp.power(Y[:, :, N] - g_tf(X[:, :, N]), 2))
    # loss += torch.sum(torch.pow(Z1 - Dg_tf(X1), 2)) #terminal 1st order condition remove for now
    return loss
Пример #19
0
def func(var=np.array([0.5, 1.])):
    temp = np.subtract(arg, var[1])
    temp = np.power(temp, 2)
    temp_sum = np.sum(temp)
    divider = np.power(var[0], 2)
    res = np.divide(temp_sum, divider)
    common = np.log(np.divide(1, var[0]))
    return np.dot(res, common)
Пример #20
0
 def for_body(i, state):
   x, first_moment, second_moment = state
   g = grad_f(x)
   first_moment = beta_1*first_moment + (1 - beta_1)*g
   second_moment = beta_2*second_moment + (1-beta_2)*(g**2)
   true_first_moment = first_moment/(1. - np.power(beta_1, i+1))
   true_second_moment = second_moment/(1. - np.power(beta_2, i+1))
   x = x - lr*true_first_moment/(np.sqrt(true_second_moment) + epsilon)
   return (x, first_moment, second_moment)
Пример #21
0
def odeInt(f, x0, y0, xT):
    # 1993 Solving Ordinary Differential Equations I, page 169
    # initial step size
    f0 = f(x0, y0)
    d0 = rmsNorm(y0)
    d1 = rmsNorm(f0)
    if d0 < 1e-5 or d1 < 1e-5:
        h0 = 1e-6
    else:
        h0 = 1e-2 * d0 / d1
    y1 = y0 + f0 * h0
    f1 = f(x0 + h0, y1)
    d2 = rmsNorm(f1 - f0) / h0
    maxD = np.maximum(d1, d2)
    if maxD <= 1e-15:
        h1 = np.maximum(1e-6, h0 * 1e-3)
    else:
        h1 = np.power(1e-2 / maxD, 1 / (P + 1))
    step = np.minimum(1e2 * h0, h1)
    # integrate
    x = x0
    y = y0
    k1 = f0
    while x < xT:
        rejected = False
        accepted = False
        while not accepted:
            ks, y1, y1h = odeIntStep(step, f, x, y, k1)
            xNew = x + step
            scale = ABS_TOL + np.maximum(np.abs(y1), np.abs(y1h)) * REL_TOL
            errNorm = rmsNorm((y1 - y1h) / scale)
            if errNorm < 1:
                accepted = True
                if errNorm == 0:
                    updateFactor = MAX_UPDATE_FACTOR
                else:
                    updateFactor = np.minimum(
                        MAX_UPDATE_FACTOR,
                        SAFETY_FACTOR * np.power(errNorm, ERROR_EXP))
                if rejected:
                    updateFactor = np.minimum(1, updateFactor)
                step *= updateFactor
            else:
                rejected = True
                updateFactor = np.maximum(
                    MIN_UPDATE_FACTOR,
                    SAFETY_FACTOR * np.power(errNorm, ERROR_EXP))
                step *= updateFactor
        # interpolate
        # update
        x = xNew
        y = y1
        k1 = ks[6]
    step = xT - x
    ks, y1, y1h = odeIntStep(step, f, x, y, k1)
    return y1
Пример #22
0
def sum_f(var):
    ret = 0
    for i in range(10**4):
        # x = 10 / 10**4 * i - 5
        x = np.subtract(np.divide(10, np.dot(np.power(10, 4), i)), 5)
        # ret -= (x - var[1])**2 / var[0]**2
        ret -= np.subtract(
            ret,
            np.divide(np.power(np.subtract(x, var[1]), 2), np.power(var[0],
                                                                    2)))
    return np.dot(np.log(np.divide(1, var[0])), ret)
Пример #23
0
def sersic(x, y, mux=0, muy=0, roll=0, q=1, c=2, I=1, Re=1, n=1):
    rm = __rotation_matrix(roll)
    qm = jnp.array(((1 / q, 0), (0, 1)))
    mu = jnp.array((mux, muy))
    P = jnp.stack((x.ravel(), y.ravel()))
    dp = (jnp.expand_dims(mu, 1) - P)
    R = jnp.power(
        jnp.sum(jnp.power(jnp.abs(jnp.dot(qm, jnp.dot(rm, dp))), c), axis=0),
        1 / c).reshape(x.shape)
    intensity = I * jnp.exp(-_b(n) * ((R / Re)**(1 / n) - 1))
    return intensity
Пример #24
0
def compute_sigma(num_examples, iterations, lipshitz, strong, diameter,
                  epsilon, delta):
    """Theorem 3.1 https://arxiv.org/pdf/2007.02923.pdf"""
    gamma = (smooth - strong) / (smooth + strong)
    numerator = 4 * np.sqrt(2) * (lipshitz + smooth * diameter) * np.power(
        gamma, iterations)
    denominator = (strong * num_examples *
                   (1 - np.power(gamma, iterations))) * (
                       (np.sqrt(np.log(1 / delta) + epsilon)) -
                       np.sqrt(np.log(1 / delta)))
    return numerator / denominator
Пример #25
0
 def model():
     with numpyro.plate("plate", 2):
         loc_latent = numpyro.sample(
             "loc_latent", FakeNormal(loc0, jnp.power(lam0, -0.5)))
         for i, x in enumerate(data):
             numpyro.sample(
                 "obs_{}".format(i),
                 dist.Normal(loc_latent, jnp.power(lam, -0.5)),
                 obs=x,
             )
     return loc_latent
Пример #26
0
def test_tracegraph_normal_normal():
    # normal-normal; known covariance
    lam0 = jnp.array([0.1, 0.1])  # precision of prior
    loc0 = jnp.array([0.0, 0.5])  # prior mean
    # known precision of observation noise
    lam = jnp.array([6.0, 4.0])
    data = []
    data.append(jnp.array([-0.1, 0.3]))
    data.append(jnp.array([0.0, 0.4]))
    data.append(jnp.array([0.2, 0.5]))
    data.append(jnp.array([0.1, 0.7]))
    n_data = len(data)
    sum_data = data[0] + data[1] + data[2] + data[3]
    analytic_lam_n = lam0 + n_data * lam
    analytic_log_sig_n = -0.5 * jnp.log(analytic_lam_n)
    analytic_loc_n = sum_data * (lam / analytic_lam_n) + loc0 * (
        lam0 / analytic_lam_n)

    class FakeNormal(dist.Normal):
        reparametrized_params = []

    def model():
        with numpyro.plate("plate", 2):
            loc_latent = numpyro.sample(
                "loc_latent", FakeNormal(loc0, jnp.power(lam0, -0.5)))
            for i, x in enumerate(data):
                numpyro.sample(
                    "obs_{}".format(i),
                    dist.Normal(loc_latent, jnp.power(lam, -0.5)),
                    obs=x,
                )
        return loc_latent

    def guide():
        loc_q = numpyro.param("loc_q",
                              analytic_loc_n + jnp.array([0.334, 0.334]))
        log_sig_q = numpyro.param(
            "log_sig_q", analytic_log_sig_n + jnp.array([-0.29, -0.29]))
        sig_q = jnp.exp(log_sig_q)
        with numpyro.plate("plate", 2):
            loc_latent = numpyro.sample("loc_latent", FakeNormal(loc_q, sig_q))
        return loc_latent

    adam = optim.Adam(step_size=0.0015, b1=0.97, b2=0.999)
    svi = SVI(model, guide, adam, loss=TraceGraph_ELBO())
    svi_result = svi.run(jax.random.PRNGKey(0), 5000)

    loc_error = jnp.sum(
        jnp.power(analytic_loc_n - svi_result.params["loc_q"], 2.0))
    log_sig_error = jnp.sum(
        jnp.power(analytic_log_sig_n - svi_result.params["log_sig_q"], 2.0))

    assert_allclose(loc_error, 0, atol=0.05)
    assert_allclose(log_sig_error, 0, atol=0.05)
Пример #27
0
def loss_function(params, t, W, Xzero):  #idea take M,D out by making X0
    loss = 0
    X, Y, Y_tilde = vXYpaths(params, t, W, Xzero)

    loss += jnp.sum(jnp.power(Y[:, :, 1:] - Y_tilde,
                              2))  # Y is 51, Y_tilde is 50

    loss += jnp.sum(jnp.power(Y[:, :, N] - g_tf(X[:, :, N]),
                              2))  #terminal condition

    # loss += torch.sum(torch.pow(Z1 - Dg_tf(X1), 2)) #terminal 1st order condition remove for now
    return loss
Пример #28
0
def harmonic_angle(conf, params, box, angle_idxs, param_idxs, cos_angles=True):
    """
    Compute the harmonic bond energy given a collection of molecules.

    This implements a harmonic angle potential: V(t) = k*(t - t0)^2 or V(t) = k*(cos(t)-cos(t0))^2

    Parameters:
    -----------
    conf: shape [num_atoms, 3] np.array
        atomic coordinates

    params: shape [num_params,] np.array
        unique parameters

    box: shape [3, 3] np.array
        periodic boundary vectors, if not None

    angle_idxs: shape [num_angles, 3] np.array
        each element (a, b, c) is a unique angle in the conformation. atom b is defined
        to be the middle atom.

    param_idxs: shape [num_angles, 2] np.array
        each element (k_idx, t_idx) maps into params for angle constants and ideal angles

    cos_angles: True (default)
        if True, then this instead implements V(t) = k*(cos(t)-cos(t0))^2. This is far more
        numerically stable when the angle is pi.

    """
    ci = conf[angle_idxs[:, 0]]
    cj = conf[angle_idxs[:, 1]]
    ck = conf[angle_idxs[:, 2]]

    kas = params[param_idxs[:, 0]]
    a0s = params[param_idxs[:, 1]]

    vij = delta_r(ci, cj, box)
    vjk = delta_r(ck, cj, box)

    top = np.sum(np.multiply(vij, vjk), -1)
    bot = np.linalg.norm(vij, axis=-1) * np.linalg.norm(vjk, axis=-1)

    tb = top / bot

    # (ytz): we used the squared version so that we make this energy being strictly positive
    if cos_angles:
        energies = kas / 2 * np.power(tb - np.cos(a0s), 2)
    else:
        angle = np.arccos(tb)
        energies = kas / 2 * np.power(angle - a0s, 2)

    return np.sum(energies, -1)  # reduce over all angles
Пример #29
0
def lens_against_ft(Uin, wvl, d1, f):
    N = len(Uin[:,0])
    k = 2 * jnp.pi / wvl
    fx = []
    x2 = []
    y2 = []
    for i in range(512):
        fx.append((-256 + i) * wvl * f)
    for i in range(512):
        x2.append(fx)
    y2 = jnp.transpose(x2)
    Uout = jnp.exp(complex("j") * k / (2*f) * (jnp.power(x2, 2) + jnp.power(y2, 2))) / (complex("j")*wvl*f) * ft2(Uin, d1)
    return Uout, np.array(x2), np.array(y2)
Пример #30
0
def func(var):
    # (var[0]-args[:,0])**3
    temp = np.sin(np.power(np.subtract(var[0], args[:, 0]), 3))
    ele1 = np.sum(temp)
    # sin(var[1]**3 - var[2]**2) * args[:,1]
    temp = np.sin(np.subtract(np.power(var[1], 3), np.power(var[2], 2)))
    ele2 = np.sum(np.dot(temp, args[:, 1]))
    # log(args[:,1]*args[:,2]*var[0]*var[1]*var[2])
    temp = np.multiply(args[:, 1], args[:, 2])
    ele3 = np.sum(np.log(1 + np.abs(np.dot(var[0] * var[1] * var[2], temp))))
    # print("\n\n", ele1, "\n", ele2, "\n", ele3)
    # print(res, type(res))
    return ele1 + ele2 + ele3