Evaluate diffusion with non-isotropic diffusivity

[WessZumino]

Hi, first of all I love the work you are doing here. I am working on an advection-diffusion model of a scalar field C; when evaluating the diffusion term with a non-isotropic diffusivity, a vector field is returned instead of a scalar. I have discussed this issue in private with @holl-, reposting here for visibility and tracking.

Consider the following setup: $\bf{D} = (D_x , D_y , D_z)$ is a vector field, where each component is space independent for now. However, the value of each term is different. I create this using:

from phi.torch import flow as pf

Diffusivity = pf.CenteredGrid(
    values=(.1, 1., 1.),
    extrapolation=pf.extrapolation.BOUNDARY,
    bounds=conc.bounds,
    resolution=conc.resolution

where $C \equiv conc$ is a real scalar field in 3D. In diffuse.explicit the following happens:

def explicit(field: FieldType,
             diffusivity: float or math.Tensor or Field,
             dt: float or math.Tensor,
             substeps: int = 1) -> FieldType:

amount = diffusivity * dt
if isinstance(amount, Field):
    amount = amount.at(field)
for i in range(substeps):
    field += amount / substeps * laplace(field).with_extrapolation(field.extrapolation)

return field

when D is a real number, this correctly evaluates: $D \nabla^2 C$ that is a scalar. However, when using D as a vector field, although this is correctly resampled at C, it evaluates: $\bf{D} \nabla^2 C$, that is a vector. The reason is that the general diffusion term is: $$\nabla \cdot (\bf{D} \nabla^2 C) = \partial^i (D_i \partial_i C)$$ When $D_i$ is constant, the expression reduces to $D^i \partial^2_i C$, that is a kind of "weighted" sum of each component of the vector field $\partial^2_i C$. When using the laplace() method in the code above, this evaluates first $\sum_i \partial^2_i C$, that is the origin of the issue.

As discussed, a simple solution here is to allow a weight term in the definition of the laplace() method, so that we could pass laplace(field, diffusivity) and obtain a scalar field now. Is there a simple way of implementing this? I tried a couple of options but I failed when trying to integrate with your codebase. I would truly appreciate your help.

Thank you

[holl-]

I believe adding a weights parameter to math.laplace() should solve this. The important line is

result = math.sum_(result, '_laplace')

which should be replaced by something like

if weights is not None:
    result *= rename_dims(weights, channel, '_laplace')
result = math.sum_(result, '_laplace')

This parameter should also be added to field.laplace() to resample the weights field as needed.

[WessZumino]

Following your suggestion, I have added:

if weights is not None:
        result *= rename_dims(weights, "vector", batch("_laplace"))
result = math. Sum(result, "_laplace")

then in field.laplace() I tried to add two arguments to the lambda function, but it did not work. Another option (probably wrong) is to use (I am assuming weights exists below):

def laplace(field: GridType, weights: GridType, axes=spatial) -> GridType:
"""Finite-difference laplace operator for Grids. See `phi.math.laplace()`."""

     result = weights._op1(
          lambda tensor_weight: field._op1(
          lambda tensor_field: math.laplace(
              x=tensor_field, dx=field.dx, weights=tensor_weight, 
             padding=field. Extrapolation, dims=axes
          )
      )
)

    return result

but this does not work either. What is the right way of passing two fields to field.laplace() ?

Thanks!

[holl-]

Hi, I got it to work like this:

def laplace(field: GridType, axes=spatial, weights: Tensor or Field = None) -> GridType:
    if isinstance(weights, Field):
        weights = weights.at(field).values
    result = field._op1(lambda tensor: math.laplace(tensor, dx=field.dx, padding=field.extrapolation, dims=axes, weights=weights))
    return result

I've also added a check to the diffuse.explicit function to check whether a channel dimension is present and interpet that as non-isotropic diffusion. I've pushed the changes into develop: c6924f14896cd7ed9ac5d80f965b57103a83bd07.