<commit_before>import os
import struct
import hashlib

import h5py
import numpy as np

from hyperion.util.meshgrid import meshgrid_nd
from hyperion.util.functions import FreezableClass, link_or_copy
from hyperion.util.logger import logger


def zero_density(grid, xmin=-np.inf, xmax=np.inf, ymin=-np.inf, ymax=np.inf, zmin=np.inf, zmax=np.inf):
    for ilevel, level in enumerate(grid.levels):
        for igrid, grid in enumerate(level.grids):
            wx = np.linspace(grid.xmin, grid.xmax, grid.nx + 1)
            wy = np.linspace(grid.ymin, grid.ymax, grid.ny + 1)
            wz = np.linspace(grid.zmin, grid.zmax, grid.nz + 1)
            x = 0.5 * (wx[:-1] + wx[1:])
            y = 0.5 * (wy[:-1] + wy[1:])
            z = 0.5 * (wz[:-1] + wz[1:])
            gx, gy, gz = meshgrid_nd(x, y, z)
            reset = (gx < xmin) | (gx > xmax) | (gy < ymin) | (gy > ymax) | (gz < zmin) | (gz > zmax)
            grid.data[reset] = 0.
    return grid


class Grid(FreezableClass):

    def __init__(self):

        # The 3D data arrays (can have various components)
        self.quantities = {}

        # The boundaries of the 3D grid in real space
        self.xmin, self.xmax = None, None
        self.ymin, self.ymax = None, None
        self.zmin, self.zmax = None, None

        # The dimensions of the array
        self.nx, self.ny, self.nz = None, None, None

        self._freeze()

    def __getattr__(self, attribute):
        if attribute == 'shape':
            return (self.nz, self.ny, self.nx)
        else:
            return FreezableClass.__getattr__(self, attribute)


class Level(FreezableClass):

    def __init__(self):

        # The list of grids in the level
        self.grids = []

        self._freeze()


class AMRGrid(FreezableClass):

    def __init__(self, amr_grid=None):

        # Initalize AMR levels
        self.levels = []

        self._freeze()

        # Copy geometry if provided
        if amr_grid is not None:
            for level in amr_grid.levels:
                level_ref = Level()
                for grid in level.grids:
                    grid_ref = Grid()
                    grid_ref.nx = grid.nx
                    grid_ref.ny = grid.ny
                    grid_ref.nz = grid.nz
                    grid_ref.xmin, grid_ref.xmax = grid.xmin, grid.xmax
                    grid_ref.ymin, grid_ref.ymax = grid.ymin, grid.ymax
                    grid_ref.zmin, grid_ref.zmax = grid.zmin, grid.zmax
                    grid_ref.quantities = {}
                    level_ref.grids.append(grid_ref)
                self.levels.append(level_ref)

    def __getattr__(self, attribute):
        if attribute == 'shape':
            return (1, 1, self.ncells)
        elif attribute == 'n_dust':
            n_dust = None
            for level in self.levels:
                for grid in level.grids:
                    for quantity in grid.quantities:
                        if type(grid.quantities[quantity]) in [list, tuple]:
                            if n_dust is None:
                                n_dust = len(grid.quantities[quantity])
                            else:
                                if n_dust != len(grid.quantities[quantity]):
                                    raise ValueError("Not all dust lists in the AMR grid have the same size")
            return n_dust
        else:
            return FreezableClass.__getattribute__(self, attribute)

    def _check_array_dimensions(self, amr_grid=None):
        '''
        Check that a grid's array dimensions agree with this grid's metadata

        Parameters
        ----------
        amr_grid: AMR grid, optional
            The AMR grid for which to test the array dimensions. If this is not
            specified, this method performs a self-consistency check of array
            dimensions and meta-data.
        '''

        # If no grid is specified, do a self-consistency checks
        if amr_grid is None:
            amr_grid = self

        # Loop over levels
        for ilevel, level_ref in enumerate(self.levels):

            # Read in level
            level = amr_grid.levels[ilevel]

            # Loop over grids
            for igrid, grid_ref in enumerate(level.grids):

                # Read in grid
                grid = level.grids[igrid]

                # Loop over quantities
                for quantity in grid.quantities:

                    # Extract array
                    array = grid.quantities[quantity]

                    if type(array) in [list, tuple]:

                        # Check that dimensions are compatible
                        for item in array:
                            if item.shape != grid_ref.shape:
                                raise ValueError("Arrays in list do not have the right "
                                                 "dimensions: %s instead of %s"
                                                 % (item.shape, grid_ref.shape))

                    elif type(array) == np.ndarray:

                        if array.shape != grid_ref.shape:
                            raise ValueError("Array does not have the right "
                                             "dimensions: %s instead of %s"
                                             % (array.shape, grid_ref.shape))

                    elif isinstance(array, h5py.ExternalLink):

                        array = h5py.File(array.filename, 'r')[array.path]

                        if len(array.shape) == 3:

                            if array.shape != grid_ref.shape:
                                raise ValueError("Array does not have the right "
                                                 "dimensions: %s instead of %s"
                                                 % (array.shape, grid_ref.shape))

                        elif len(array.shape) == 4:

                            for item in array:
                                if item.shape != grid_ref.shape:
                                    raise ValueError("Arrays in list do not have the right "
                                                     "dimensions: %s instead of %s"
                                                     % (item.shape, grid_ref.shape))

                        else:

                            raise Exception("Unexpected number of dimensions: %i" % array.ndim)

                    else:

                        raise ValueError("Array should be a list or a Numpy array")

    def read(self, group, quantities='all'):
        '''
        Read in an AMR grid

        Parameters
        ----------
        group: h5py.Group
            The HDF5 group to read the grid from
        quantities: 'all' or list
            Which physical quantities to read in. Use 'all' to read in all
            quantities or a list of strings to read only specific quantities.
        '''

        # Initialize levels
        self.levels = []

        # Read geometry group
        g_geometry = group['Geometry']

        # Read quantities group
        g_quantities = group['Quantities']

        # Check that grid is indeed AMR
        if g_geometry.attrs['grid_type'] != 'amr':
            raise Exception("Grid is not AMR")
<commit_msg>Updated AMRGrid to follow same API as regular 3D grids<commit_after>import os
import struct
import hashlib

import h5py
import numpy as np

from hyperion.util.meshgrid import meshgrid_nd
from hyperion.util.functions import FreezableClass, link_or_copy
from hyperion.util.logger import logger
from hyperion.grid.grid_helpers import single_grid_dims


def zero_density(grid, xmin=-np.inf, xmax=np.inf, ymin=-np.inf, ymax=np.inf, zmin=np.inf, zmax=np.inf):
    for ilevel, level in enumerate(grid.levels):
        for igrid, grid in enumerate(level.grids):
            wx = np.linspace(grid.xmin, grid.xmax, grid.nx + 1)
            wy = np.linspace(grid.ymin, grid.ymax, grid.ny + 1)
            wz = np.linspace(grid.zmin, grid.zmax, grid.nz + 1)
            x = 0.5 * (wx[:-1] + wx[1:])
            y = 0.5 * (wy[:-1] + wy[1:])
            z = 0.5 * (wz[:-1] + wz[1:])
            gx, gy, gz = meshgrid_nd(x, y, z)
            reset = (gx < xmin) | (gx > xmax) | (gy < ymin) | (gy > ymax) | (gz < zmin) | (gz > zmax)
            grid.data[reset] = 0.
    return grid


class Grid(FreezableClass):

    def __init__(self):

        # The 3D data arrays (can have various components)
        self.quantities = {}

        # The boundaries of the 3D grid in real space
        self.xmin, self.xmax = None, None
        self.ymin, self.ymax = None, None
        self.zmin, self.zmax = None, None

        # The dimensions of the array
        self.nx, self.ny, self.nz = None, None, None

        self._freeze()

    def __getattr__(self, attribute):
        if attribute == 'shape':
            return (self.nz, self.ny, self.nx)
        else:
            return FreezableClass.__getattr__(self, attribute)


class Level(FreezableClass):

    def __init__(self):

        # The list of grids in the level
        self.grids = []

        self._freeze()


class AMRGrid(FreezableClass):

    def __init__(self, amr_grid=None):

        # Initalize AMR levels
        self.levels = []

        self._freeze()

        # Copy geometry if provided
        if amr_grid is not None:
            for level in amr_grid.levels:
                level_ref = Level()
                for grid in level.grids:
                    grid_ref = Grid()
                    grid_ref.nx = grid.nx
                    grid_ref.ny = grid.ny
                    grid_ref.nz = grid.nz
                    grid_ref.xmin, grid_ref.xmax = grid.xmin, grid.xmax
                    grid_ref.ymin, grid_ref.ymax = grid.ymin, grid.ymax
                    grid_ref.zmin, grid_ref.zmax = grid.zmin, grid.zmax
                    grid_ref.quantities = {}
                    level_ref.grids.append(grid_ref)
                self.levels.append(level_ref)

    def __getattr__(self, attribute):
        if attribute == 'shape':
            return (1, 1, self.ncells)
        elif attribute == 'n_dust':
            n_dust = None
            for level in self.levels:
                for grid in level.grids:
                    for quantity in grid.quantities:
                        n_dust_q, shape_q = single_grid_dims(grid.quantities[quantity])
                        if n_dust is None:
                            n_dust = n_dust_q
                        else:
                            if n_dust != n_dust_q:
                                raise ValueError("Not all dust lists in the grid have the same size")
            return n_dust
        else:
            return FreezableClass.__getattribute__(self, attribute)

    def _check_array_dimensions(self, amr_grid=None):
        '''
        Check that a grid's array dimensions agree with this grid's metadata

        Parameters
        ----------
        amr_grid: AMR grid, optional
            The AMR grid for which to test the array dimensions. If this is not
            specified, this method performs a self-consistency check of array
            dimensions and meta-data.
        '''

        # If no grid is specified, do a self-consistency checks
        if amr_grid is None:
            amr_grid = self

        # Loop over levels
        for ilevel, level_ref in enumerate(self.levels):

            # Read in level
            level = amr_grid.levels[ilevel]

            # Loop over grids
            for igrid, grid_ref in enumerate(level.grids):

                # Read in grid
                grid = level.grids[igrid]

                # Loop over quantities
                for quantity in grid.quantities:

                    n_pop, shape = single_grid_dims(grid.quantities[quantity])

                    if shape != grid.shape:
                        raise ValueError("Quantity arrays do not have the right "
                                         "dimensions: %s instead of %s"
                                         % (shape, grid.shape))

    def read(self, group, quantities='all'):
        '''
        Read in an AMR grid

        Parameters
        ----------
        group: h5py.Group
            The HDF5 group to read the grid from
        quantities: 'all' or list
            Which physical quantities to read in. Use 'all' to read in all
            quantities or a list of strings to read only specific quantities.
        '''

        # Initialize levels
        self.levels = []

        # Read geometry group
        g_geometry = group['Geometry']

        # Read quantities group
        g_quantities = group['Quantities']

        # Check that grid is indeed AMR
        if g_geometry.attrs['grid_type'] != 'amr':
            raise Exception("Grid is not AMR")

{'fraction_non_alphanumeric': 0.08742355339501333, 'fraction_numerical': 0.003998745491610475, 'mean_word_length': 1.4212224753227032, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 24, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6580247', 'n_tokens_mistral': 3631, 'n_tokens_neox': 3402, 'n_words': 1051}