Python numpy.empty_like() Examples

def elastic_transform(image, alpha=1000, sigma=30, spline_order=1, mode='nearest', random_state=np.random):
    """Elastic deformation of image as described in [Simard2003]_.
    .. [Simard2003] Simard, Steinkraus and Platt, "Best Practices for
       Convolutional Neural Networks applied to Visual Document Analysis", in
       Proc. of the International Conference on Document Analysis and
       Recognition, 2003.
    """
    assert image.ndim == 3
    shape = image.shape[:2]

    dx = gaussian_filter((random_state.rand(*shape) * 2 - 1),
                         sigma, mode="constant", cval=0) * alpha
    dy = gaussian_filter((random_state.rand(*shape) * 2 - 1),
                         sigma, mode="constant", cval=0) * alpha

    x, y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing='ij')
    indices = [np.reshape(x + dx, (-1, 1)), np.reshape(y + dy, (-1, 1))]
    result = np.empty_like(image)
    for i in range(image.shape[2]):
        result[:, :, i] = map_coordinates(
            image[:, :, i], indices, order=spline_order, mode=mode).reshape(shape)
    return result 

def canonicalize(mf, mo_coeff_kpts, mo_occ_kpts, fock=None):
    if fock is None:
        dm = mf.make_rdm1(mo_coeff_kpts, mo_occ_kpts)
        fock = mf.get_fock(dm=dm)
    mo_coeff = []
    mo_energy = []
    for k, mo in enumerate(mo_coeff_kpts):
        mo1 = np.empty_like(mo)
        mo_e = np.empty_like(mo_occ_kpts[k])
        occidx = mo_occ_kpts[k] == 2
        viridx = ~occidx
        for idx in (occidx, viridx):
            if np.count_nonzero(idx) > 0:
                orb = mo[:,idx]
                f1 = reduce(np.dot, (orb.T.conj(), fock[k], orb))
                e, c = scipy.linalg.eigh(f1)
                mo1[:,idx] = np.dot(orb, c)
                mo_e[idx] = e
        mo_coeff.append(mo1)
        mo_energy.append(mo_e)
    return mo_energy, mo_coeff 

def level_thickness(p, p_top=0., p_bot=1.01325e5):
    """
    Calculates the thickness, in Pa, of each pressure level.

    Assumes that the pressure values given are at the center of that model
    level, except for the lowest value (typically 1000 hPa), which is the
    bottom boundary. The uppermost level extends to 0 hPa.

    Unlike `dp_from_p`, this does not incorporate the surface pressure.

    """
    p_vals = to_pascal(p.values.copy())
    dp_vals = np.empty_like(p_vals)
    # Bottom level extends from p[0] to halfway betwen p[0] and p[1].
    dp_vals[0] = p_bot - 0.5*(p_vals[0] + p_vals[1])
    # Middle levels extend from halfway between [k-1], [k] and [k], [k+1].
    dp_vals[1:-1] = 0.5*(p_vals[0:-2] - p_vals[2:])
    # Top level extends from halfway between top two levels to 0 hPa.
    dp_vals[-1] = 0.5*(p_vals[-2] + p_vals[-1]) - p_top
    dp = p.copy()
    dp.values = dp_vals
    return dp 

def finish(self):
        """ Compute the aggregated value in each scenario based on the parent `EventRecorder` events """
        events = self.event_recorder.events
        # Return NaN if no events found
        if len(events) == 0:
            return

        scen_id = np.empty(len(events), dtype=np.int)
        values = np.empty_like(scen_id, dtype=np.float64)

        for i, evt in enumerate(events):
            scen_id[i] = evt.scenario_index.global_id
            values[i] = pandas.Series(evt.values).aggregate(self.event_agg_func)

        df = pandas.DataFrame({'scenario_id': scen_id, 'value': values})

        # Group by scenario ...
        grouped = df.groupby('scenario_id').agg(self.recorder_agg_func)
        # ... and update the internal values
        for index, row in grouped.iterrows():
            self._values[index] = row['value'] 

def canonicalize(mf, mo_coeff, mo_occ, fock=None):
    '''Canonicalization diagonalizes the Fock matrix within occupied, open,
    virtual subspaces separatedly (without change occupancy).
    '''
    if fock is None:
        dm = mf.make_rdm1(mo_coeff, mo_occ)
        fock = mf.get_fock(dm=dm)
    coreidx = mo_occ == 2
    viridx = mo_occ == 0
    openidx = ~(coreidx | viridx)
    mo = numpy.empty_like(mo_coeff)
    mo_e = numpy.empty(mo_occ.size)
    for idx in (coreidx, openidx, viridx):
        if numpy.count_nonzero(idx) > 0:
            orb = mo_coeff[:,idx]
            f1 = reduce(numpy.dot, (orb.conj().T, fock, orb))
            e, c = scipy.linalg.eigh(f1)
            mo[:,idx] = numpy.dot(orb, c)
            mo_e[idx] = e
    return mo_e, mo 

def time_reversal_matrix(mol, mat):
    ''' T(A_ij) = A[T(i),T(j)]^*
    '''
    n2c = mol.nao_2c()
    tao = numpy.asarray(mol.time_reversal_map())
    # tao(i) = -j  means  T(f_i) = -f_j
    # tao(i) =  j  means  T(f_i) =  f_j
    idx = abs(tao)-1 # -1 for C indexing convention
    #:signL = [(1 if x>0 else -1) for x in tao]
    #:sign = numpy.hstack((signL, signL))

    #:tmat = numpy.empty_like(mat)
    #:for j in range(mat.__len__()):
    #:    for i in range(mat.__len__()):
    #:        tmat[idx[i],idx[j]] = mat[i,j] * sign[i]*sign[j]
    #:return tmat.conjugate()
    sign_mask = tao<0
    if mat.shape[0] == n2c*2:
        idx = numpy.hstack((idx, idx+n2c))
        sign_mask = numpy.hstack((sign_mask, sign_mask))

    tmat = mat.take(idx,axis=0).take(idx,axis=1)
    tmat[sign_mask,:] *= -1
    tmat[:,sign_mask] *= -1
    return tmat.T 

def contract_1e(f1e, fcivec, norb, nelec, link_index=None):
    fcivec = numpy.asarray(fcivec, order='C')
    link_index = _unpack(norb, nelec, link_index)
    na, nlink = link_index.shape[:2]
    assert(fcivec.size == na**2)
    ci1 = numpy.empty_like(fcivec)
    f1e_tril = lib.pack_tril(f1e)
    libfci.FCIcontract_1e_spin0(f1e_tril.ctypes.data_as(ctypes.c_void_p),
                                fcivec.ctypes.data_as(ctypes.c_void_p),
                                ci1.ctypes.data_as(ctypes.c_void_p),
                                ctypes.c_int(norb), ctypes.c_int(na),
                                ctypes.c_int(nlink),
                                link_index.ctypes.data_as(ctypes.c_void_p))
# no *.5 because FCIcontract_2e_spin0 only compute half of the contraction
    return lib.transpose_sum(ci1, inplace=True).reshape(fcivec.shape)

# Note eri is NOT the 2e hamiltonian matrix, the 2e hamiltonian is
# h2e = eri_{pq,rs} p^+ q r^+ s
#     = (pq|rs) p^+ r^+ s q - (pq|rs) \delta_{qr} p^+ s
# so eri is defined as
#       eri_{pq,rs} = (pq|rs) - (1/Nelec) \sum_q (pq|qs)
# to restore the symmetry between pq and rs,
#       eri_{pq,rs} = (pq|rs) - (.5/Nelec) [\sum_q (pq|qs) + \sum_p (pq|rp)]
# Please refer to the treatment in direct_spin1.absorb_h1e
# the input fcivec should be symmetrized 

def reorder_eri(eri, norb, orbsym):
    if orbsym is None:
        return [eri], numpy.arange(norb), numpy.zeros(norb,dtype=numpy.int32)
# map irrep IDs of Dooh or Coov to D2h, C2v
# see symm.basis.linearmole_symm_descent
    orbsym = numpy.asarray(orbsym) % 10
# irrep of (ij| pair
    trilirrep = (orbsym[:,None]^orbsym)[numpy.tril_indices(norb)]
# and the number of occurence for each irrep
    dimirrep = numpy.asarray(numpy.bincount(trilirrep), dtype=numpy.int32)
# we sort the irreps of (ij| pair, to group the pairs which have same irreps
# "order" is irrep-id-sorted index. The (ij| paired is ordered that the
# pair-id given by order[0] comes first in the sorted pair
# "rank" is a sorted "order". Given nth (ij| pair, it returns the place(rank)
# of the sorted pair
    old_eri_irrep = numpy.asarray(trilirrep, dtype=numpy.int32)
    rank_in_irrep = numpy.empty_like(old_eri_irrep)
    p0 = 0
    eri_irs = [numpy.zeros((0,0))] * TOTIRREPS
    for ir, nnorb in enumerate(dimirrep):
        idx = numpy.asarray(numpy.where(trilirrep == ir)[0], dtype=numpy.int32)
        rank_in_irrep[idx] = numpy.arange(nnorb, dtype=numpy.int32)
        eri_irs[ir] = lib.take_2d(eri, idx, idx)
        p0 += nnorb
    return eri_irs, rank_in_irrep, old_eri_irrep 

def test_type2_rad_part(self):
        rc = .8712
        rs, ws = radi.gauss_chebyshev(99)
        def type2_facs_rad(ish, lc):
            facs0 = facs_rad(mol, ish, lc, rc, rs).transpose(0,2,1).copy()
            facs1 = numpy.empty_like(facs0)
            libecp.type2_facs_rad(facs1.ctypes.data_as(ctypes.c_void_p),
                                  ctypes.c_int(ish), ctypes.c_int(lc),
                                  ctypes.c_double(rc),
                                  rs.ctypes.data_as(ctypes.c_void_p),
                                  ctypes.c_int(len(rs)), ctypes.c_int(1),
                                  mol._atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.natm),
                                  mol._bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.nbas),
                                  mol._env.ctypes.data_as(ctypes.c_void_p))
            self.assertTrue(numpy.allclose(facs0, facs1))
        for ish in range(mol.nbas):
            for lc in range(5):
                type2_facs_rad(ish, lc) 

def gen_vind(mf, mo_coeff, mo_occ):
    nao, nmo = mo_coeff.shape
    mocc = mo_coeff[:,mo_occ>0]
    nocc = mocc.shape[1]
    vresp = mf.gen_response(mo_coeff, mo_occ, hermi=1)
    def fx(mo1):
        mo1 = mo1.reshape(-1,nmo,nocc)
        nset = len(mo1)
        dm1 = numpy.empty((nset,nao,nao))
        for i, x in enumerate(mo1):
            dm = reduce(numpy.dot, (mo_coeff, x*2, mocc.T)) # *2 for double occupancy
            dm1[i] = dm + dm.T
        v1 = vresp(dm1)
        v1vo = numpy.empty_like(mo1)
        for i, x in enumerate(v1):
            v1vo[i] = reduce(numpy.dot, (mo_coeff.T, x, mocc))
        return v1vo
    return fx 

def _solve_mo1_uncoupled(mo_energy, mo_occ, h1, s1):
    '''uncoupled first order equation'''
    e_a = mo_energy[mo_occ==0]
    e_i = mo_energy[mo_occ>0]
    e_ai = 1 / (e_a.reshape(-1,1) - e_i)

    hs = h1 - s1 * e_i

    mo10 = numpy.empty_like(hs)
    mo10[:,mo_occ==0,:] = -hs[:,mo_occ==0,:] * e_ai
    mo10[:,mo_occ>0,:] = -s1[:,mo_occ>0,:] * .5

    e_ji = e_i.reshape(-1,1) - e_i
    mo_e10 = hs[:,mo_occ>0,:] + mo10[:,mo_occ>0,:] * e_ji
    return mo10, mo_e10

#TODO: merge to hessian.rhf.solve_mo1 function 

def _add_giao_phase(mol, vmat):
    '''Add the factor i/2*(Ri-Rj) of the GIAO phase e^{i/2 (Ri-Rj) times r}'''
    ao_coords = rhf_mag._get_ao_coords(mol)
    Rx = .5 * (ao_coords[:,0:1] - ao_coords[:,0])
    Ry = .5 * (ao_coords[:,1:2] - ao_coords[:,1])
    Rz = .5 * (ao_coords[:,2:3] - ao_coords[:,2])
    vxc20 = numpy.empty_like(vmat)
    vxc20[0]  = Ry * vmat[2] - Rz * vmat[1]
    vxc20[1]  = Rz * vmat[0] - Rx * vmat[2]
    vxc20[2]  = Rx * vmat[1] - Ry * vmat[0]
    vxc20, vmat = vmat, vxc20
    vxc20[:,0]  = Ry * vmat[:,2] - Rz * vmat[:,1]
    vxc20[:,1]  = Rz * vmat[:,0] - Rx * vmat[:,2]
    vxc20[:,2]  = Rx * vmat[:,1] - Ry * vmat[:,0]
    vxc20 *= -1
    return vxc20 

def mask_randomly(self, imgs):
        y1 = np.random.randint(0, self.img_rows - self.mask_height, imgs.shape[0])
        y2 = y1 + self.mask_height
        x1 = np.random.randint(0, self.img_rows - self.mask_width, imgs.shape[0])
        x2 = x1 + self.mask_width

        masked_imgs = np.empty_like(imgs)
        missing_parts = np.empty((imgs.shape[0], self.mask_height, self.mask_width, self.channels))
        for i, img in enumerate(imgs):
            masked_img = img.copy()
            _y1, _y2, _x1, _x2 = y1[i], y2[i], x1[i], x2[i]
            missing_parts[i] = masked_img[_y1:_y2, _x1:_x2, :].copy()
            masked_img[_y1:_y2, _x1:_x2, :] = 0
            masked_imgs[i] = masked_img

        return masked_imgs, missing_parts, (y1, y2, x1, x2) 

def log_loss_value(Z, weights, total_weights, rho):
    """
    computes the value and slope of the logistic loss in a numerically stable way
    supports sample non-negative weights for each example in the training data
    see http://stackoverflow.com/questions/20085768/

    Parameters
    ----------
    Z               numpy.array containing training data with shape = (n_rows, n_cols)
    rho             numpy.array of coefficients with shape = (n_cols,)
    total_weights   numpy.sum(total_weights) (only included to reduce computation)
    weights         numpy.array of sample weights with shape (n_rows,)

    Returns
    -------
    loss_value  scalar = 1/n_rows * sum(log( 1 .+ exp(-Z*rho))

    """
    scores = Z.dot(rho)
    pos_idx = scores > 0
    loss_value = np.empty_like(scores)
    loss_value[pos_idx] = np.log1p(np.exp(-scores[pos_idx]))
    loss_value[~pos_idx] = -scores[~pos_idx] + np.log1p(np.exp(scores[~pos_idx]))
    loss_value = loss_value.dot(weights) / total_weights
    return loss_value 

def log_loss_value(Z, rho):
    """
    computes the value and slope of the logistic loss in a numerically stable way
    see also: http://stackoverflow.com/questions/20085768/

    Parameters
    ----------
    Z           numpy.array containing training data with shape = (n_rows, n_cols)
    rho         numpy.array of coefficients with shape = (n_cols,)

    Returns
    -------
    loss_value  scalar = 1/n_rows * sum(log( 1 .+ exp(-Z*rho))

    """
    scores = Z.dot(rho)
    pos_idx = scores > 0
    loss_value = np.empty_like(scores)
    loss_value[pos_idx] = np.log1p(np.exp(-scores[pos_idx]))
    loss_value[~pos_idx] = -scores[~pos_idx] + np.log1p(np.exp(scores[~pos_idx]))
    loss_value = loss_value.mean()
    return loss_value 

def log_probs(Z, rho):
    """
    compute the probabilities of the logistic loss function in a way that is numerically stable

    see also: http://stackoverflow.com/questions/20085768/
    Parameters
    ----------
    Z           numpy.array containing training data with shape = (n_rows, n_cols)
    rho         numpy.array of coefficients with shape = (n_cols,)

    Returns
    -------
    log_probs   numpy.array of probabilities under the logit model
    """

    scores = Z.dot(rho)
    pos_idx = scores > 0
    log_probs = np.empty_like(scores)
    log_probs[pos_idx]  = 1.0 / (1.0 + np.exp(-scores[pos_idx]))
    log_probs[~pos_idx] = np.exp(scores[~pos_idx]) / (1.0 + np.exp(scores[~pos_idx]))
    return log_probs 

def maximum_filter_2d(arr2D, footprint): ## Make sure arr2D is our datatype float32 and footprint of int32
    arr2DMaxed = numpy.empty_like(arr2D)
    head, tail = os.path.split(os.path.abspath(__file__)) # Used so that we can always get the kernel which should be in the same directory as this file

    maxFunction = open(head + "/2DSlidingMaxFootprintKernel.c", "rt")
    maxFunction = SourceModule(maxFunction.read())
    slidingMaxKernel = maxFunction.get_function("slidingMaxiumum2D")

    blockSize = [16, 16] # To-do: Add a variable to this, can affect performance based on GPU
    gridSize = getGridSize(blockSize, arr2D.shape) # Get the size of our grid based on the size of a grid (blocksize)


    slidingMaxKernel(cuda.In(arr2D),                   # Input
                    cuda.Out(arr2DMaxed),              # Output
                    numpy.int32(footprint.shape[1]),   # Kernel Size
                    numpy.int32(arr2D.shape[1]),       # Row Stride
                    numpy.int32(1),                    # Column Stride
                    numpy.int32(int(arr2D.shape[1])),  # Array Column Count
                    numpy.int32(int(arr2D.shape[0])),  # Array Row Count
                    cuda.In(footprint),
                    block=(blockSize[0],blockSize[1],1),
                    grid=(gridSize[0],gridSize[1],1)
    )

    return arr2DMaxed 

def _bounds_from_array(arr, dim_name, bounds_name):
    """Get the bounds of an array given its center values.

    E.g. if lat-lon grid center lat/lon values are known, but not the
    bounds of each grid box.  The algorithm assumes that the bounds
    are simply halfway between each pair of center values.
    """
    # TODO: don't assume needed dimension is in axis=0
    # TODO: refactor to get rid of repetitive code
    spacing = arr.diff(dim_name).values
    lower = xr.DataArray(np.empty_like(arr), dims=arr.dims,
                         coords=arr.coords)
    lower.values[:-1] = arr.values[:-1] - 0.5*spacing
    lower.values[-1] = arr.values[-1] - 0.5*spacing[-1]
    upper = xr.DataArray(np.empty_like(arr), dims=arr.dims,
                         coords=arr.coords)
    upper.values[:-1] = arr.values[:-1] + 0.5*spacing
    upper.values[-1] = arr.values[-1] + 0.5*spacing[-1]
    bounds = xr.concat([lower, upper], dim='bounds')
    return bounds.T 

def test_cholesky_eri(self):
        j2c = df.incore.fill_2c2e(mol, auxmol)
        eri0 = numpy.empty_like(j2c)
        pi = 0
        for i in range(mol.nbas, len(bas)):
            pj = 0
            for j in range(mol.nbas, len(bas)):
                shls = (i, j)
                buf = gto.moleintor.getints_by_shell('int2c2e_sph',
                                                     shls, atm, bas, env)
                di, dj = buf.shape
                eri0[pi:pi+di,pj:pj+dj] = buf
                pj += dj
            pi += di
        self.assertTrue(numpy.allclose(eri0, j2c))

        j3c = df.incore.aux_e2(mol, auxmol, intor='int3c2e_sph', aosym='s2ij')
        cderi = df.incore.cholesky_eri(mol)
        eri0 = numpy.einsum('pi,pk->ik', cderi, cderi)
        eri1 = numpy.einsum('ik,kl->il', j3c, numpy.linalg.inv(j2c))
        eri1 = numpy.einsum('ip,kp->ik', eri1, j3c)
        self.assertTrue(numpy.allclose(eri1, eri0))

        cderi1 = df.incore.cholesky_eri_debug(mol)
        self.assertAlmostEqual(abs(cderi-cderi1).max(), 0, 9) 

def check_synced(self):
        if self.comm.Get_rank() == 0: # this is root
            theta = self.getflat()
            self.comm.Bcast(theta, root=0)
        else:
            thetalocal = self.getflat()
            thetaroot = np.empty_like(thetalocal)
            self.comm.Bcast(thetaroot, root=0)
            assert (thetaroot == thetalocal).all(), (thetaroot, thetalocal) 

def log_loss_value_from_scores(scores):
    """
    computes the logistic loss value from a vector of scores in a numerically stable way
    where scores = Z.dot(rho)

    see also: http://stackoverflow.com/questions/20085768/

    this function is used for heuristics (discrete_descent, sequential_rounding).
    to save computation when running the heuristics, we store the scores and
    call this function to compute the loss directly from the scores
    this reduces the need to recompute the dot product.

    Parameters
    ----------
    scores  numpy.array of scores = Z.dot(rho)

    Returns
    -------
    loss_value  scalar = 1/n_rows * sum(log( 1 .+ exp(-Z*rho))

    """

    pos_idx = scores > 0
    loss_value = np.empty_like(scores)
    loss_value[pos_idx] = np.log1p(np.exp(-scores[pos_idx]))
    loss_value[~pos_idx] = -scores[~pos_idx] + np.log1p(np.exp(scores[~pos_idx]))
    loss_value = loss_value.mean()
    return loss_value 

def reset(self):
        """Resets the iterator to the beginning of the data."""
        self.curr_idx = 0
        random.shuffle(self.idx)
        for buck in self.data:
            np.random.shuffle(buck)

        self.nddata = []
        self.ndlabel = []
        for buck in self.data:
            label = np.empty_like(buck)
            label[:, :-1] = buck[:, 1:]
            label[:, -1] = self.invalid_label
            self.nddata.append(ndarray.array(buck, dtype=self.dtype))
            self.ndlabel.append(ndarray.array(label, dtype=self.dtype)) 

def check_synced(self):
        if self.comm.Get_rank() == 0: # this is root
            theta = self.getflat()
            self.comm.Bcast(theta, root=0)
        else:
            thetalocal = self.getflat()
            thetaroot = np.empty_like(thetalocal)
            self.comm.Bcast(thetaroot, root=0)
            assert (thetaroot == thetalocal).all(), (thetaroot, thetalocal) 

def log_loss_value_from_scores(weights, total_weights, scores):
    """
    computes the logistic loss value from a vector of scores in a numerically stable way
    where scores = Z.dot(rho)

    see also: http://stackoverflow.com/questions/20085768/

    this function is used for heuristics (discrete_descent, sequential_rounding).
    to save computation when running the heuristics, we store the scores and
    call this function to compute the loss directly from the scores
    this reduces the need to recompute the dot product.

    Parameters
    ----------
    scores          numpy.array of scores = Z.dot(rho)
    total_weights   numpy.sum(total_weights) (only included to reduce computation)
    weights         numpy.array of sample weights with shape (n_rows,)

    Returns
    -------
    loss_value  scalar = 1/n_rows * sum(log( 1 .+ exp(-Z*rho))

    """
    pos_idx = scores > 0
    loss_value = np.empty_like(scores)
    loss_value[pos_idx] = np.log1p(np.exp(-scores[pos_idx]))
    loss_value[~pos_idx] = -scores[~pos_idx] + np.log1p(np.exp(scores[~pos_idx]))
    loss_value = loss_value.dot(weights) / total_weights

    return loss_value 

def spherical_to_cart(x):
  shape = x.shape
  x = x.reshape(-1,3)
  y = np.empty_like(x)
  y[:,0] = x[:,0] * np.sin(x[:,1]) * np.cos(x[:,2])
  y[:,1] = x[:,0] * np.sin(x[:,1]) * np.sin(x[:,2])
  y[:,2] = x[:,0] * np.cos(x[:,1])
  return y.reshape(shape) 

def mesh_triangle_areas(verts, faces):
  a = verts[faces[:,0]]
  b = verts[faces[:,1]]
  c = verts[faces[:,2]]
  x = np.empty_like(a)
  x = a - b
  y = a - c
  t = np.empty_like(a)
  t[:,0] = (x[:,1] * y[:,2] - x[:,2] * y[:,1]);
  t[:,1] = (x[:,2] * y[:,0] - x[:,0] * y[:,2]);
  t[:,2] = (x[:,0] * y[:,1] - x[:,1] * y[:,0]);
  return np.linalg.norm(t, axis=1) / 2 

def get_mesh(rng, min_z=0):
  # set up background board
  verts, faces, normals, colors = [], [], [], []
  v, f, n = co.geometry.xyplane(z=0, interleaved=True)
  v[:,2] += -v[:,2].min() + rng.uniform(2,7)
  v[:,:2] *= 5e2
  v[:,2] = np.mean(v[:,2]) + (v[:,2] - np.mean(v[:,2])) * 5e2
  c = np.empty_like(v)
  c[:] = rng.uniform(0,1, size=(3,)).astype(np.float32)
  verts.append(v)
  faces.append(f)
  normals.append(n)
  colors.append(c)

  # randomly sample 4 foreground objects for each scene
  for shape_idx in range(4):
    v, f, n = objs[rng.randint(0,len(objs))]
    v, f, n = v.copy(), f.copy(), n.copy()

    s = rng.uniform(0.25, 1)
    v *= s
    R = co.geometry.rotm_from_quat(co.geometry.quat_random(rng=rng))
    v = v @ R.T
    n = n @ R.T
    v[:,2] += -v[:,2].min() + min_z + rng.uniform(0.5, 3)
    v[:,:2] += rng.uniform(-1, 1, size=(1,2))

    c = np.empty_like(v)
    c[:] = rng.uniform(0,1, size=(3,)).astype(np.float32)

    verts.append(v.astype(np.float32))
    faces.append(f)
    normals.append(n)
    colors.append(c)

  verts, faces = co.geometry.stack_mesh(verts, faces)
  normals = np.vstack(normals).astype(np.float32)
  colors = np.vstack(colors).astype(np.float32)
  return verts, faces, colors, normals 

def forward(self, Z: np.ndarray) -> np.ndarray:
        output = np.empty_like(Z)
        p = Z > 0.
        n = ~p
        output[p] = Z[p]
        output[n] = self.alpha * Z[n]
        return output 

def check_synced(self):
        if self.comm.Get_rank() == 0: # this is root
            theta = self.getflat()
            self.comm.Bcast(theta, root=0)
        else:
            thetalocal = self.getflat()
            thetaroot = np.empty_like(thetalocal)
            self.comm.Bcast(thetaroot, root=0)
            assert (thetaroot == thetalocal).all(), (thetaroot, thetalocal) 

def cart_to_spherical(x):
  shape = x.shape
  x = x.reshape(-1,3)
  y = np.empty_like(x)
  y[:,0] = np.linalg.norm(x, axis=1)  # r
  y[:,1] = np.arccos(x[:,2] / y[:,0]) # theta
  y[:,2] = np.arctan2(x[:,1], x[:,0]) # phi
  return y.reshape(shape)