Use SymPy to generate the 3D rotation matrix and products for diffusion_stencil_3d.
349882
def _symbolic_rotation_helper():
    """"""
    from sympy import symbols, Matrix

    cpsi, spsi = symbols('cpsi, spsi')
    cth, sth = symbols('cth, sth')
    cphi, sphi = symbols('cphi, sphi')
    Rpsi = Matrix([[cpsi, spsi, 0], [-spsi, cpsi, 0], [0, 0, 1]])
    Rth = Matrix([[1, 0, 0], [0, cth, sth], [0, -sth, cth]])
    Rphi = Matrix([[cphi, sphi, 0], [-sphi, cphi, 0], [0, 0, 1]])

    Q = Rpsi * Rth * Rphi

    epsy, epsz = symbols('epsy, epsz')
    A = Matrix([[1, 0, 0], [0, epsy, 0], [0, 0, epsz]])

    D = Q * A * Q.T

    for i in range(3):
        for j in range(3):
            print('D[%d, %d] = %s' % (i, j, D[i, j]))
600535
def Ry_matrix(theta):
    """"""
    return np.array([
        [np.cos(theta), 0, np.sin(theta)],
        [0, 1, 0],
        [-np.sin(theta), 0, np.cos(theta)]
    ])
677639
def symmetrise(matrix, tri='upper'):
    
    if tri == 'upper':
        tri_fn = np.triu_indices
    else:
        tri_fn = np.tril_indices
    size = matrix.shape[0]
    matrix[tri_fn(size)[::-1]] = matrix[tri_fn(size)]
    return matrix
89714
def de_projection_3d(amplitudes, sigmas):
    
    amplitudes_3d = amplitudes / sigmas / np.sqrt(2*np.pi)
    return amplitudes_3d, sigmas
33747
def rotated(self, angle_degrees_ccw):
        """"""
        angle = angle_degrees_ccw / 180.0 * pi
        c, s = cos(angle), sin(angle)
        return self @ PdfMatrix((c, s, -s, c, 0, 0))
17299
def spm(x, model, y=None, n_samples=None, dx_min=-0.1,
        dx_max=0.1, n_dxs=5, dy_min=-0.1, dy_max=0.1, n_dys=5,
        angle_min=-30, angle_max=30, n_angles=31, black_border_size=0):
  
  if y is None:
    preds = model.get_probs(x)
    # Using model predictions as ground truth to avoid label leaking
    preds_max = reduce_max(preds, 1, keepdims=True)
    y = tf.to_float(tf.equal(preds, preds_max))
    y = tf.stop_gradient(y)
    del preds
  y = y / reduce_sum(y, 1, keepdims=True)

  # Define the range of transformations
  dxs = np.linspace(dx_min, dx_max, n_dxs)
  dys = np.linspace(dy_min, dy_max, n_dys)
  angles = np.linspace(angle_min, angle_max, n_angles)

  if n_samples is None:
    import itertools
    transforms = list(itertools.product(*[dxs, dys, angles]))
  else:
    sampled_dxs = np.random.choice(dxs, n_samples)
    sampled_dys = np.random.choice(dys, n_samples)
    sampled_angles = np.random.choice(angles, n_samples)
    transforms = zip(sampled_dxs, sampled_dys, sampled_angles)
  transformed_ims = parallel_apply_transformations(
      x, transforms, black_border_size)

  def _compute_xent(x):
    preds = model.get_logits(x)
    return tf.nn.softmax_cross_entropy_with_logits_v2(
        labels=y, logits=preds)

  all_xents = tf.map_fn(
      _compute_xent,
      transformed_ims,
      parallel_iterations=1)  # Must be 1 to avoid keras race conditions

  # Return the adv_x with worst accuracy

  # all_xents is n_total_samples x batch_size (SB)
  all_xents = tf.stack(all_xents)  # SB

  # We want the worst case sample, with the largest xent_loss
  worst_sample_idx = tf.argmax(all_xents, axis=0)  # B

  batch_size = tf.shape(x)[0]
  keys = tf.stack([
      tf.range(batch_size, dtype=tf.int32),
      tf.cast(worst_sample_idx, tf.int32)
  ], axis=1)
  transformed_ims_bshwc = tf.einsum('sbhwc->bshwc', transformed_ims)
  after_lookup = tf.gather_nd(transformed_ims_bshwc, keys)  # BHWC
  return after_lookup
1200458
def get_symmetry_operations(self, cartesian=False):
        
        rotation, translation = self._get_symmetry()
        symmops = []
        mat = self._structure.lattice.matrix.T
        invmat = np.linalg.inv(mat)
        for rot, trans in zip(rotation, translation):
            if cartesian:
                rot = np.dot(mat, np.dot(rot, invmat))
                trans = np.dot(trans, self._structure.lattice.matrix)
            op = SymmOp.from_rotation_and_translation(rot, trans)
            symmops.append(op)
        return symmops
989122
def G(w, Xs):
    ''''''
    n = len(Xs)
    P = projection_matrix(w)
    Ys = [np.dot(P, X) for X in Xs]
    A = calc_A(Ys)
    A_hat = calc_A_hat(A, skew_matrix(w))

    
    u = sum(np.dot(Y, Y) for Y in Ys) / n
    v = np.dot(A_hat, sum(np.dot(Y, Y) * Y for Y in Ys)) / np.trace(np.dot(A_hat, A))

    return sum((np.dot(Y, Y) - u - 2 * np.dot(Y, v)) ** 2 for Y in Ys)
1074996
def get_rotations(self):
        """"""
        if self.centrosymmetric:
            return np.vstack((self.rotations, -self.rotations))
        else:
            return self.rotations
1257925
def iu1transform(self, zeta):
        """"""
        x = real(zeta)
        y = imag(zeta)
        z = x*self.radius.real + y*self.radius.imag
        return self.rot_matrix*z + self.center
756032
def symmetrize(A):
    
    assert type(A) is spmatrix, "argument must be a sparse matrix"
    assert A.size[0] == A.size[1], "argument must me a square matrix"
    idx = [i for i,ij in enumerate(zip(A.I,A.J)) if ij[0] > ij[1]]
    return tril(A) + spmatrix(A.V[idx], A.J[idx], A.I[idx], A.size)
27838
def _derive_stereographic():
    """"""
    from sympy import symbols, atan2, acos, rot_axis1, rot_axis3, Matrix
    x_c, y_c, z_c, x, y, z = symbols('x_c y_c z_c x y z')

    # The angles we'll need to rotate through.
    around_z = atan2(x_c, y_c)
    around_x = acos(-z_c)

    # Apply rotations to produce an "o" = output vector.
    v = Matrix([x, y, z])
    xo, yo, zo = rot_axis1(around_x) * rot_axis3(-around_z) * v

    # Which we then use the stereographic projection to produce the
    # final "p" = plotting coordinates.
    xp = xo / (1 - zo)
    yp = yo / (1 - zo)

    return xp, yp
565436
def _build_matrix(p, q, deriv):
    """"""
    A = [([1 for _ in range(-p, q+1)])]
    for i in range(1, p + q + 1):
        A.append([j**i for j in range(-p, q+1)])
    return np.array(A)
1131398
def cross_product_matrix(vec):
    """"""
    return np.array([[0, -vec[2], vec[1]],
                     [vec[2], 0, -vec[0]],
                     [-vec[1], vec[0], 0]])
129669
def symbolic_Rz_matrix(symbolic_theta):
    """"""
    return sympy.Matrix([
        [sympy.cos(symbolic_theta), -sympy.sin(symbolic_theta), 0],
        [sympy.sin(symbolic_theta), sympy.cos(symbolic_theta), 0],
        [0, 0, 1]
    ])
710856
def isosurface_from_data(data, isolevel, origin, spacing):
    
    spacing = np.array(extent/resolution)
    if isolevel >= 0:
        triangles = marching_cubes(data, isolevel)
    else: # Wrong traingle unwinding roder -- god only knows why
        triangles = marching_cubes(-data, -isolevel)
    faces = []
    verts = []
    for i, t in enumerate(triangles):
       faces.append([i * 3, i * 3 +1, i * 3 + 2])
       verts.extend(t)
       
    faces = np.array(faces)
    verts = origin + spacing/2 + np.array(verts)*spacing
    return verts, faces
470058
def get_matrix(theta, phi, psi, dx, dy, dz,
                   matrix=np.zeros((4, 4), dtype=DTYPE),
                   angles=np.zeros(3, dtype=DTYPE)):
        
        # NB!: matrix and angles by default are being overwritten on each call
        # thus, only created once at compile time.

        angles[0] = theta
        angles[1] = phi
        angles[2] = psi

        cx, cy, cz = np.cos(angles)
        sx, sy, sz = np.sin(angles)

        rotation = matrix[:3, :3]
        rotation.flat = (cx * cz - sx * cy * sz,
                         cx * sz + sx * cy * cz, sx * sy,
                         -sx * cz - cx * cy * sz,
                         -sx * sz + cx * cy * cz, cx * sy,
                         sy * sz,
                         -sy * cz, cy)

        # Translation component
        matrix[:3, 3] = dx, dy, dz
        matrix[3, 3] = 1.
        return matrix