Exemplo n.º 1
0
def sparse_cg(field, A, max_iterations, guess, accuracy, back_prop=False):
    div_vec = math.reshape(field, [-1, int(np.prod(field.shape[1:]))])
    if guess is not None:
        guess = math.reshape(guess, [-1, int(np.prod(field.shape[1:]))])
    apply_A = lambda pressure: math.matmul(A, pressure)
    result_vec, iterations = conjugate_gradient(div_vec, apply_A, guess, accuracy, max_iterations, back_prop)
    return math.reshape(result_vec, math.shape(field)), iterations
Exemplo n.º 2
0
    def solve(self, field, domain, guess, enable_backprop):
        assert isinstance(domain, FluidDomain)
        active_mask = domain.active_tensor(extend=1)
        fluid_mask = domain.accessible_tensor(extend=1)
        dimensions = math.staticshape(field)[1:-1]
        N = int(np.prod(dimensions))
        periodic = Material.periodic(domain.domain.boundaries)

        if math.choose_backend([field, active_mask,
                                fluid_mask]).matches_name('SciPy'):
            A = sparse_pressure_matrix(dimensions, active_mask, fluid_mask,
                                       periodic)
        else:
            sidx, sorting = sparse_indices(dimensions, periodic)
            sval_data = sparse_values(dimensions, active_mask, fluid_mask,
                                      sorting, periodic)
            backend = math.choose_backend(field)
            sval_data = backend.cast(sval_data, field.dtype)
            A = backend.sparse_tensor(indices=sidx,
                                      values=sval_data,
                                      shape=[N, N])

        div_vec = math.reshape(field, [-1, int(np.prod(field.shape[1:]))])
        if guess is not None:
            guess = math.reshape(guess, [-1, int(np.prod(field.shape[1:]))])

        def apply_A(pressure):
            return math.matmul(A, pressure)

        result_vec, iterations = conjugate_gradient(div_vec, apply_A, guess,
                                                    self.accuracy,
                                                    self.max_iterations,
                                                    enable_backprop)
        return math.reshape(result_vec, math.shape(field)), iterations
Exemplo n.º 3
0
    def multi_advect(self, fields, interpolation="LINEAR", dt=1):
        assert isinstance(
            fields,
            (list, tuple)), "first parameter must be either a tuple or list"
        inputs_lists = []
        coords_lists = []
        value_generators = []
        for field in fields:
            if isinstance(field, StaggeredGrid):
                i, c, v = self._mac_block_advection(field.staggered, dt)
            else:
                i, c, v = self._centered_block_advection(field, dt)
            inputs_lists.append(i)
            coords_lists.append(c)
            value_generators.append(v)

        inputs = math.concat(sum(inputs_lists, []), 0)
        coords = math.concat(sum(coords_lists, []), 0)
        all_advected = math.resample(inputs,
                                     coords,
                                     interpolation=interpolation,
                                     boundary="REPLICATE")
        all_advected = math.reshape(all_advected, [self.spatial_rank, -1] +
                                    list(all_advected.shape[1:]))
        all_advected = math.unstack(all_advected)
        results = []
        abs_i = 0
        for i in range(len(inputs_lists)):
            n = len(inputs_lists[0])
            assigned_advected = all_advected[abs_i:abs_i + n]
            results.append(value_generators[i](assigned_advected))
            abs_i += n
        return results
Exemplo n.º 4
0
def upsample2x(tensor, interpolation="LINEAR"):
    if interpolation.lower() != "linear":
        raise ValueError("Only linear interpolation supported")
    dims = range(spatial_rank(tensor))
    vlen = tensor.shape[-1]
    spatial_dims = tensor.shape[1:-1]
    tensor = math.pad(tensor,
                      [[0, 0]] + [[1, 1]] * spatial_rank(tensor) + [[0, 0]],
                      "SYMMETRIC")
    for dim in dims:
        left_slices_1 = [(slice(2, None) if i == dim else slice(None))
                         for i in dims]
        left_slices_2 = [(slice(1, -1) if i == dim else slice(None))
                         for i in dims]
        right_slices_1 = [(slice(1, -1) if i == dim else slice(None))
                          for i in dims]
        right_slices_2 = [(slice(-2) if i == dim else slice(None))
                          for i in dims]
        left = 0.75 * tensor[[slice(None)] + left_slices_2 +
                             [slice(None)]] + 0.25 * tensor[
                                 [slice(None)] + left_slices_1 + [slice(None)]]
        right = 0.25 * tensor[[slice(None)] + right_slices_2 + [
            slice(None)
        ]] + 0.75 * tensor[[slice(None)] + right_slices_1 + [slice(None)]]
        combined = math.stack([right, left], axis=2 + dim)
        tensor = math.reshape(combined, [-1] + [
            spatial_dims[dim] * 2 if i == dim else tensor.shape[i + 1]
            for i in dims
        ] + [vlen])
    return tensor
Exemplo n.º 5
0
def batch_indices(indices):
    """
Reshapes the indices such that, aside from indices, they also contain batch number.
For example the entry (32, 40) as coordinates of batch 2 will become (2, 32, 40).
Transform shape (b, p, d) to (b, p, d+1) where batch size is b, number of particles is p and number of dimensions is d. 
    """
    batch_size = indices.shape[0]
    out_spatial_rank = len(indices.shape) - 2
    out_spatial_size = math.shape(indices)[1:-1]

    batch_range = math.DYNAMIC_BACKEND.choose_backend(indices).range(batch_size)
    batch_ids = math.reshape(batch_range, [batch_size] + [1] * out_spatial_rank)
    tile_shape = math.pad(out_spatial_size, [[1,0]], constant_values=1)
    batch_ids = math.expand_dims(math.tile(batch_ids, tile_shape), axis=-1)

    return math.concat((batch_ids, indices), axis=-1)