def build(self, mesh_instance, size, g2r_field):
field_dict = {}
for key, attr in self.attr_dict.items():
if isinstance(attr.dtype, CompoundType):
field_dict[key] = attr.dtype.field(shape=None,
needs_grad=attr.needs_grad)
else:
field_dict[key] = impl.field(attr.dtype,
shape=None,
needs_grad=attr.needs_grad)
if self.layout == Layout.SOA:
for key, field in field_dict.items():
impl.root.dense(impl.axes(0), size).place(field)
if self.attr_dict[key].needs_grad:
impl.root.dense(impl.axes(0), size).place(field.grad)
elif len(field_dict) > 0:
impl.root.dense(impl.axes(0),
size).place(*tuple(field_dict.values()))
grads = []
for key, field in field_dict.items():
if self.attr_dict[key].needs_grad: grads.append(field.grad)
if len(grads) > 0:
impl.root.dense(impl.axes(0), size).place(*grads)
return MeshElementField(mesh_instance, self._type, self.attr_dict,
field_dict, g2r_field)
def get_normals_field(vertices):
if vertices not in normals_field_cache:
N = vertices.shape[0]
normals = Vector.field(3, f32, shape=(N, ))
normal_weights = field(f32, shape=(N, ))
normals_field_cache[vertices] = (normals, normal_weights)
return (normals, normal_weights)
return normals_field_cache[vertices]
def field(cls,
members,
shape=None,
name="<Struct>",
offset=None,
needs_grad=False,
layout=Layout.AOS):
if shape is None and offset is not None:
raise TaichiSyntaxError(
"shape cannot be None when offset is being set")
field_dict = {}
for key, dtype in members.items():
field_name = name + '.' + key
if isinstance(dtype, CompoundType):
field_dict[key] = dtype.field(shape=None,
name=field_name,
offset=offset,
needs_grad=needs_grad)
else:
field_dict[key] = impl.field(dtype,
shape=None,
name=field_name,
offset=offset,
needs_grad=needs_grad)
if shape is not None:
if isinstance(shape, numbers.Number):
shape = (shape, )
if isinstance(offset, numbers.Number):
offset = (offset, )
if offset is not None and len(shape) != len(offset):
raise TaichiSyntaxError(
f'The dimensionality of shape and offset must be the same ({len(shape)} != {len(offset)})'
)
dim = len(shape)
if layout == Layout.SOA:
for e in field_dict.values():
ti.root.dense(impl.index_nd(dim),
shape).place(e, offset=offset)
if needs_grad:
for e in field_dict.values():
ti.root.dense(impl.index_nd(dim),
shape).place(e.grad, offset=offset)
else:
ti.root.dense(impl.index_nd(dim),
shape).place(*tuple(field_dict.values()),
offset=offset)
if needs_grad:
grads = tuple(e.grad for e in field_dict.values())
ti.root.dense(impl.index_nd(dim),
shape).place(*grads, offset=offset)
return StructField(field_dict, name=name)
def allocate_dual(x, dual_root):
"""Allocate dual field for forward mode autodiff
"""
dtype = x.dtype
shape = x.shape
dim = len(shape)
x_dual = impl.field(dtype)
x._set_grad(x_dual, reverse_mode=False)
x._get_field_members()[0].ptr.set_dual(x_dual._get_field_members()[0].ptr)
dual_root.dense(impl.index_nd(dim), shape).place(x_dual)
def field(cls,
members,
shape=None,
name="<Struct>",
offset=None,
needs_grad=False,
layout=Layout.AOS):
"""Creates a :class:`~taichi.StructField` with each element
has this struct as its type.
Args:
members (dict): a dict, each item is like `name: type`.
shape (Tuple[int]): width and height of the field.
offset (Tuple[int]): offset of the indices of the created field.
For example if `offset=(-10, -10)` the indices of the field
will start at `(-10, -10)`, not `(0, 0)`.
needs_grad (bool): enabling gradient field or not.
layout: AOS or SOA.
Example:
>>> vec3 = ti.types.vector(3, ti.f32)
>>> sphere = {"center": vec3, "radius": float}
>>> F = ti.Struct.field(sphere, shape=(3, 3))
>>> F
{'center': array([[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]],
[[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]]], dtype=float32), 'radius': array([[0., 0., 0.],
[0., 0., 0.],
[0., 0., 0.]], dtype=float32)}
"""
if shape is None and offset is not None:
raise TaichiSyntaxError(
"shape cannot be None when offset is being set")
field_dict = {}
for key, dtype in members.items():
field_name = name + '.' + key
if isinstance(dtype, CompoundType):
field_dict[key] = dtype.field(shape=None,
name=field_name,
offset=offset,
needs_grad=needs_grad)
else:
field_dict[key] = impl.field(dtype,
shape=None,
name=field_name,
offset=offset,
needs_grad=needs_grad)
if shape is not None:
if isinstance(shape, numbers.Number):
shape = (shape, )
if isinstance(offset, numbers.Number):
offset = (offset, )
if offset is not None and len(shape) != len(offset):
raise TaichiSyntaxError(
f'The dimensionality of shape and offset must be the same ({len(shape)} != {len(offset)})'
)
dim = len(shape)
if layout == Layout.SOA:
for e in field_dict.values():
impl.root.dense(impl.index_nd(dim),
shape).place(e, offset=offset)
if needs_grad:
for e in field_dict.values():
impl.root.dense(impl.index_nd(dim),
shape).place(e.grad, offset=offset)
else:
impl.root.dense(impl.index_nd(dim),
shape).place(*tuple(field_dict.values()),
offset=offset)
if needs_grad:
grads = tuple(e.grad for e in field_dict.values())
impl.root.dense(impl.index_nd(dim),
shape).place(*grads, offset=offset)
return StructField(field_dict, name=name)
def __init__(self, filename):
with open(filename, "r") as fi:
data = json.loads(fi.read())
self.num_patches = data["num_patches"]
self.element_fields = {}
self.relation_fields = {}
self.num_elements = {}
self.max_num_per_patch = {}
for element in data["elements"]:
element_type = MeshElementType(element["order"])
self.num_elements[element_type] = element["num"]
self.max_num_per_patch[element_type] = element["max_num_per_patch"]
element["l2g_mapping"] = np.array(element["l2g_mapping"])
element["l2r_mapping"] = np.array(element["l2r_mapping"])
element["g2r_mapping"] = np.array(element["g2r_mapping"])
self.element_fields[element_type] = {}
self.element_fields[element_type]["owned"] = impl.field(
dtype=ti.i32, shape=self.num_patches + 1)
self.element_fields[element_type]["total"] = impl.field(
dtype=ti.i32, shape=self.num_patches + 1)
self.element_fields[element_type]["l2g"] = impl.field(
dtype=ti.i32, shape=element["l2g_mapping"].shape[0])
self.element_fields[element_type]["l2r"] = impl.field(
dtype=ti.i32, shape=element["l2r_mapping"].shape[0])
self.element_fields[element_type]["g2r"] = impl.field(
dtype=ti.i32, shape=element["g2r_mapping"].shape[0])
for relation in data["relations"]:
from_order = relation["from_order"]
to_order = relation["to_order"]
rel_type = MeshRelationType(
relation_by_orders(from_order, to_order))
self.relation_fields[rel_type] = {}
self.relation_fields[rel_type]["value"] = impl.field(
dtype=ti.i32, shape=len(relation["value"]))
if from_order <= to_order:
self.relation_fields[rel_type]["offset"] = impl.field(
dtype=ti.i32, shape=len(relation["offset"]))
for element in data["elements"]:
element_type = MeshElementType(element["order"])
self.element_fields[element_type]["owned"].from_numpy(
np.array(element["owned_offsets"]))
self.element_fields[element_type]["total"].from_numpy(
np.array(element["total_offsets"]))
self.element_fields[element_type]["l2g"].from_numpy(
element["l2g_mapping"])
self.element_fields[element_type]["l2r"].from_numpy(
element["l2r_mapping"])
self.element_fields[element_type]["g2r"].from_numpy(
element["g2r_mapping"])
for relation in data["relations"]:
from_order = relation["from_order"]
to_order = relation["to_order"]
rel_type = MeshRelationType(
relation_by_orders(from_order, to_order))
self.relation_fields[rel_type]["value"].from_numpy(
np.array(relation["value"]))
if from_order <= to_order:
self.relation_fields[rel_type]["offset"].from_numpy(
np.array(relation["offset"]))
self.attrs = {}
self.attrs["x"] = np.array(data["attrs"]["x"]).reshape(-1, 3)
def field(cls,
n,
m,
dtype,
shape=None,
name="",
offset=None,
needs_grad=False,
layout=None): # TODO(archibate): deprecate layout
'''ti.Matrix.field'''
self = cls.empty(n, m)
self.entries = []
self.n = n
self.m = m
self.dt = dtype
if isinstance(dtype, (list, tuple, np.ndarray)):
# set different dtype for each element in Matrix
# see #2135
if m == 1:
assert len(np.shape(dtype)) == 1 and len(
dtype
) == n, f'Please set correct dtype list for Vector. The shape of dtype list should be ({n}, ) instead of {np.shape(dtype)}'
for i in range(n):
self.entries.append(impl.field(dtype[i], name=name))
else:
assert len(np.shape(dtype)) == 2 and len(dtype) == n and len(
dtype[0]
) == m, f'Please set correct dtype list for Matrix. The shape of dtype list should be ({n}, {m}) instead of {np.shape(dtype)}'
for i in range(n):
for j in range(m):
self.entries.append(impl.field(dtype[i][j], name=name))
else:
for _ in range(n * m):
self.entries.append(impl.field(dtype, name=name))
self.grad = self.make_grad()
if layout is not None:
assert shape is not None, 'layout is useless without shape'
if shape is None:
assert offset is None, "shape cannot be None when offset is being set"
if shape is not None:
if isinstance(shape, numbers.Number):
shape = (shape, )
if isinstance(offset, numbers.Number):
offset = (offset, )
if offset is not None:
assert len(shape) == len(
offset
), f'The dimensionality of shape and offset must be the same ({len(shape)} != {len(offset)})'
if layout is None:
layout = ti.AOS
dim = len(shape)
if layout.soa:
for i, e in enumerate(self.entries):
ti.root.dense(impl.index_nd(dim),
shape).place(e, offset=offset)
if needs_grad:
for i, e in enumerate(self.entries):
ti.root.dense(impl.index_nd(dim),
shape).place(e.grad, offset=offset)
else:
var_list = []
for i, e in enumerate(self.entries):
var_list.append(e)
ti.root.dense(impl.index_nd(dim),
shape).place(*tuple(var_list), offset=offset)
grad_var_list = []
if needs_grad:
for i, e in enumerate(self.entries):
grad_var_list.append(e.grad)
ti.root.dense(impl.index_nd(dim),
shape).place(*tuple(grad_var_list),
offset=offset)
return self
def field(cls,
n,
m,
dtype,
shape=None,
name="",
offset=None,
needs_grad=False,
layout=None): # TODO(archibate): deprecate layout
"""Construct a data container to hold all elements of the Matrix.
Args:
n (int): The desired number of rows of the Matrix.
m (int): The desired number of columns of the Matrix.
dtype (DataType, optional): The desired data type of the Matrix.
shape (Union[int, tuple of int], optional): The desired shape of the Matrix.
name (string, optional): The custom name of the field.
offset (Union[int, tuple of int], optional): The coordinate offset of all elements in a field.
needs_grad (bool, optional): Whether the Matrix need gradients.
layout (:class:`~taichi.lang.impl.Layout`, optional): The field layout, i.e., Array Of Structure(AOS) or Structure Of Array(SOA).
Returns:
:class:`~taichi.lang.matrix.Matrix`: A :class:`~taichi.lang.matrix.Matrix` instance serves as the data container.
"""
self = cls.empty(n, m)
self.entries = []
self.n = n
self.m = m
self.dt = dtype
if isinstance(dtype, (list, tuple, np.ndarray)):
# set different dtype for each element in Matrix
# see #2135
if m == 1:
assert len(np.shape(dtype)) == 1 and len(
dtype
) == n, f'Please set correct dtype list for Vector. The shape of dtype list should be ({n}, ) instead of {np.shape(dtype)}'
for i in range(n):
self.entries.append(impl.field(dtype[i], name=name))
else:
assert len(np.shape(dtype)) == 2 and len(dtype) == n and len(
dtype[0]
) == m, f'Please set correct dtype list for Matrix. The shape of dtype list should be ({n}, {m}) instead of {np.shape(dtype)}'
for i in range(n):
for j in range(m):
self.entries.append(impl.field(dtype[i][j], name=name))
else:
for _ in range(n * m):
self.entries.append(impl.field(dtype, name=name))
self.grad = self.make_grad()
if layout is not None:
assert shape is not None, 'layout is useless without shape'
if shape is None:
assert offset is None, "shape cannot be None when offset is being set"
if shape is not None:
if isinstance(shape, numbers.Number):
shape = (shape, )
if isinstance(offset, numbers.Number):
offset = (offset, )
if offset is not None:
assert len(shape) == len(
offset
), f'The dimensionality of shape and offset must be the same ({len(shape)} != {len(offset)})'
if layout is None:
layout = ti.AOS
dim = len(shape)
if layout.soa:
for i, e in enumerate(self.entries):
ti.root.dense(impl.index_nd(dim),
shape).place(e, offset=offset)
if needs_grad:
for i, e in enumerate(self.entries):
ti.root.dense(impl.index_nd(dim),
shape).place(e.grad, offset=offset)
else:
var_list = []
for i, e in enumerate(self.entries):
var_list.append(e)
ti.root.dense(impl.index_nd(dim),
shape).place(*tuple(var_list), offset=offset)
grad_var_list = []
if needs_grad:
for i, e in enumerate(self.entries):
grad_var_list.append(e.grad)
ti.root.dense(impl.index_nd(dim),
shape).place(*tuple(grad_var_list),
offset=offset)
return self
def __init__(self, data):
self.num_patches = data["num_patches"]
self.element_fields = {}
self.relation_fields = {}
self.num_elements = {}
self.max_num_per_patch = {}
for element in data["elements"]:
element_type = MeshElementType(element["order"])
self.num_elements[element_type] = element["num"]
self.max_num_per_patch[element_type] = element["max_num_per_patch"]
element["l2g_mapping"] = np.array(element["l2g_mapping"])
element["l2r_mapping"] = np.array(element["l2r_mapping"])
element["g2r_mapping"] = np.array(element["g2r_mapping"])
self.element_fields[element_type] = {}
self.element_fields[element_type]["owned"] = impl.field(
dtype=u32, shape=self.num_patches + 1)
self.element_fields[element_type]["total"] = impl.field(
dtype=u32, shape=self.num_patches + 1)
self.element_fields[element_type]["l2g"] = impl.field(
dtype=u32, shape=element["l2g_mapping"].shape[0])
self.element_fields[element_type]["l2r"] = impl.field(
dtype=u32, shape=element["l2r_mapping"].shape[0])
self.element_fields[element_type]["g2r"] = impl.field(
dtype=u32, shape=element["g2r_mapping"].shape[0])
for relation in data["relations"]:
from_order = relation["from_order"]
to_order = relation["to_order"]
rel_type = MeshRelationType(
relation_by_orders(from_order, to_order))
self.relation_fields[rel_type] = {}
self.relation_fields[rel_type]["value"] = impl.field(
dtype=u16, shape=len(relation["value"]))
if from_order <= to_order:
self.relation_fields[rel_type]["offset"] = impl.field(
dtype=u16, shape=len(relation["offset"]))
self.relation_fields[rel_type]["patch_offset"] = impl.field(
dtype=u32, shape=len(relation["patch_offset"]))
self.relation_fields[rel_type]["from_order"] = from_order
self.relation_fields[rel_type]["to_order"] = to_order
for element in data["elements"]:
element_type = MeshElementType(element["order"])
self.element_fields[element_type]["owned"].from_numpy(
np.array(element["owned_offsets"]))
self.element_fields[element_type]["total"].from_numpy(
np.array(element["total_offsets"]))
self.element_fields[element_type]["l2g"].from_numpy(
element["l2g_mapping"])
self.element_fields[element_type]["l2r"].from_numpy(
element["l2r_mapping"])
self.element_fields[element_type]["g2r"].from_numpy(
element["g2r_mapping"])
for relation in data["relations"]:
from_order = relation["from_order"]
to_order = relation["to_order"]
rel_type = MeshRelationType(
relation_by_orders(from_order, to_order))
self.relation_fields[rel_type]["value"].from_numpy(
np.array(relation["value"]))
if from_order <= to_order:
self.relation_fields[rel_type]["patch_offset"].from_numpy(
np.array(relation["patch_offset"]))
self.relation_fields[rel_type]["offset"].from_numpy(
np.array(relation["offset"]))
self.attrs = {}
self.attrs["x"] = np.array(data["attrs"]["x"]).reshape(-1, 3)
if "patcher" in data:
self.patcher = data["patcher"]
else:
self.patcher = None
def fun(arr, dtype):
field = impl.field(dtype=dtype, shape=arr.shape)
field.from_numpy(arr)
return field
