Source code for yt.data_objects.level_sets.contour_finder

from collections import defaultdict

import numpy as np

from yt.funcs import get_pbar, mylog
from yt.utilities.lib.contour_finding import (
    ContourTree,
    TileContourTree,
    link_node_contours,
    update_joins,
)
from yt.utilities.lib.partitioned_grid import PartitionedGrid


[docs] def identify_contours(data_source, field, min_val, max_val, cached_fields=None): tree = ContourTree() gct = TileContourTree(min_val, max_val) total_contours = 0 contours = {} node_ids = [] DLE = data_source.ds.domain_left_edge masks = {g.id: m for g, m in data_source.blocks} for g, node, (sl, dims, gi) in data_source.tiles.slice_traverse(): g.field_parameters.update(data_source.field_parameters) node.node_ind = len(node_ids) nid = node.node_id node_ids.append(nid) values = g[field][sl].astype("float64") contour_ids = np.zeros(dims, "int64") - 1 mask = masks[g.id][sl].astype("uint8") total_contours += gct.identify_contours( values, contour_ids, mask, total_contours ) new_contours = tree.cull_candidates(contour_ids) tree.add_contours(new_contours) # Now we can create a partitioned grid with the contours. LE = (DLE + g.dds * gi).in_units("code_length").ndarray_view() RE = LE + (dims * g.dds).in_units("code_length").ndarray_view() pg = PartitionedGrid( g.id, [contour_ids.view("float64")], mask, LE, RE, dims.astype("int64") ) contours[nid] = (g.Level, node.node_ind, pg, sl) node_ids = np.array(node_ids, dtype="int64") if node_ids.size == 0: return 0, {} trunk = data_source.tiles.tree.trunk mylog.info("Linking node (%s) contours.", len(contours)) link_node_contours(trunk, contours, tree, node_ids) mylog.info("Linked.") # joins = tree.cull_joins(bt) # tree.add_joins(joins) joins = tree.export() contour_ids = defaultdict(list) pbar = get_pbar("Updating joins ... ", len(contours)) final_joins = np.unique(joins[:, 1]) for i, nid in enumerate(sorted(contours)): level, node_ind, pg, sl = contours[nid] ff = pg.my_data[0].view("int64") update_joins(joins, ff, final_joins) contour_ids[pg.parent_grid_id].append((sl, ff)) pbar.update(i + 1) pbar.finish() rv = {} rv.update(contour_ids) # NOTE: Because joins can appear in both a "final join" and a subsequent # "join", we can't know for sure how many unique joins there are without # checking if no cells match or doing an expensive operation checking for # the unique set of final join values. return final_joins.size, rv