jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _impl_v1(cls, inputs, attr, params):
        axis = attr.get("axis", 1)
        in_shape = infer_shape(inputs[0])
        ndim = len(in_shape)
        if axis < 0:
            axis += ndim
        if axis == 0:
            reshape_shape = [-1]
        elif axis == ndim - 1:
            return _op.nn.softmax(inputs[0], axis=axis)
        else:
            axis_val = [in_shape[i] for i in range(axis)]
            reshape_shape = [np.prod(axis_val)] + [-1]
        data_reshape = _op.reshape(inputs[0], newshape=reshape_shape)
        out = _op.nn.softmax(data_reshape, axis=-1)
        out = _op.reshape(out, newshape=in_shape)
        return out
//  [python]
def white_matter_response_tournier07(
        acquisition_scheme, data, N_candidate_voxels=300, **kwargs):
    """The original white matter response estimation algorithm according to
    [1]_. In essence, it just takes the 300 voxels with the highest FA, aligns
    them with the z-axis, and estimates the averaged white matter response from
    that.

    Parameters
    ----------
    acquisition_scheme : DmipyAcquisitionScheme instance,
        An acquisition scheme that has been instantiated using dMipy.
    data : NDarray,
        Measured diffusion signal array.

    Returns
    -------
    S0_wm : positive float,
        Estimated S0 tissue response value.
    TR2_wm_model : Dmipy Anisotropic ModelFree Model
        ModelFree representation of white matter response.
    selected_indices : array of size (N_candidate_voxels,),
        indices of selected voxels for white matter response.

    References
    ----------
    .. [1] Tournier, J-Donald, Fernando Calamante, and Alan Connelly. "Robust
        determination of the fibre orientation distribution in diffusion MRI:
        non-negativity constrained super-resolved spherical deconvolution."
        Neuroimage 35.4 (2007): 1459-1472.
    """
    data_shape = np.atleast_2d(data).shape
    N_voxels = int(np.multiply.reduce(data_shape[:-1]))
    if N_voxels < N_candidate_voxels:
        msg = "The original algorithm uses 300 candidate voxels to estimate "
        msg += "the tissue response. Currently only {} ".format(N_voxels)
        msg += "candidate voxels given."
        print(msg)
        N_candidate_voxels = N_voxels

    if data.ndim == 4:
        # calculate brain mask on 4D data (x, y, z, DWI)
        b0_mask, mask = median_otsu(
            input_volume=data,
            vol_idx=np.where(acquisition_scheme.b0_mask)[0],
            median_radius=4, numpass=4)  # based on dipy default
        # needs to be eroded 3 times.
        mask_eroded = binary_erosion(mask, iterations=3)
        data_to_fit = data[mask_eroded]
    else:
        # can't calculate brain mask on other than 4D data.
        # assume the data was prepared.
        data_to_fit = data.reshape([-1, data_shape[-1]])

    gtab = gtab_dmipy2dipy(acquisition_scheme)
    tenmod = dti.TensorModel(gtab)
    tenfit = tenmod.fit(data_to_fit)
    fa = tenfit.fa

    # selected based on FA
    selected_indices = np.argsort(fa)[-N_candidate_voxels:]
    selected_data = data_to_fit[selected_indices]
    S0_wm, TR2_wm_model = estimate_TR2_anisotropic_tissue_response_model(
        acquisition_scheme, selected_data)
    return S0_wm, TR2_wm_model, selected_indices
//  [python]
def white_matter_response_tournier07(
        acquisition_scheme, data, N_candidate_voxels=300, **kwargs):
    """The original white matter response estimation algorithm according to
    [1]_. In essence, it just takes the 300 voxels with the highest FA, aligns
    them with the z-axis, and estimates the averaged white matter response from
    that.

    Parameters
    ----------
    acquisition_scheme : DmipyAcquisitionScheme instance,
        An acquisition scheme that has been instantiated using dMipy.
    data : NDarray,
        Measured diffusion signal array.

    Returns
    -------
    S0_wm : positive float,
        Estimated S0 tissue response value.
    TR2_wm_model : Dmipy Anisotropic ModelFree Model
        ModelFree representation of white matter response.
    selected_indices : array of size (N_candidate_voxels,),
        indices of selected voxels for white matter response.

    References
    ----------
    .. [1] Tournier, J-Donald, Fernando Calamante, and Alan Connelly. "Robust
        determination of the fibre orientation distribution in diffusion MRI:
        non-negativity constrained super-resolved spherical deconvolution."
        Neuroimage 35.4 (2007): 1459-1472.
    """
    data_shape = np.atleast_2d(data).shape
    N_voxels = int(np.prod(data_shape[:-1]))
    if N_voxels < N_candidate_voxels:
        msg = "The original algorithm uses 300 candidate voxels to estimate "
        msg += "the tissue response. Currently only {} ".format(N_voxels)
        msg += "candidate voxels given."
        print(msg)
        N_candidate_voxels = N_voxels

    if data.ndim == 4:
        # calculate brain mask on 4D data (x, y, z, DWI)
        b0_mask, mask = median_otsu(
            input_volume=data,
            vol_idx=np.where(acquisition_scheme.b0_mask)[0],
            median_radius=4, numpass=4)  # based on dipy default
        # needs to be eroded 3 times.
        mask_eroded = binary_erosion(mask, iterations=3)
        data_to_fit = data[mask_eroded]
    else:
        # can't calculate brain mask on other than 4D data.
        # assume the data was prepared.
        data_to_fit = data.reshape([-1, data_shape[-1]])

    gtab = gtab_dmipy2dipy(acquisition_scheme)
    tenmod = dti.TensorModel(gtab)
    tenfit = tenmod.fit(data_to_fit)
    fa = tenfit.fa

    # selected based on FA
    selected_indices = np.argsort(fa)[-N_candidate_voxels:]
    selected_data = data_to_fit[selected_indices]
    S0_wm, TR2_wm_model = estimate_TR2_anisotropic_tissue_response_model(
        acquisition_scheme, selected_data)
    return S0_wm, TR2_wm_model, selected_indices
//  [python]
def normalize(self, img, mini=0.0, maxi=1.0, copy=True, out=None):
        """Scale the intensity of the image so that the minimum is 0 and the
        maximum is 1.0 (or any value suggested).

        :param img: numpy array or pyopencl array of dim 2 or 3 and of type float
        :param mini: Expected minimum value
        :param maxi: expected maxiumum value
        :param copy: set to False to use directly the input buffer
        :param out: provides an output buffer. prevents a copy D->H

        This uses a min/max reduction in two stages plus a map operation
        """
        assert img.shape == self.buffer_shape
        events = []
        with self.sem:
            input_array, output_array = self._get_in_out_buffers(img, copy, out)
            size = numpy.int32(numpy.multiply.reduce(self.shape))
            if self.wg_red == 1:
                #  Probably on MacOS CPU WG==1 --> serial code.
                kernel = self.kernels.get_kernel("max_min_serial")
                evt = kernel(
                    self.queue,
                    (1,),
                    (1,),
                    input_array.data,
                    size,
                    self.cl_mem["max_min_d"].data,
                )
                ed = EventDescription("max_min_serial", evt)
                events.append(ed)
            else:
                stage1 = self.kernels.max_min_reduction_stage1
                stage2 = self.kernels.max_min_reduction_stage2
                local_mem = pyopencl.LocalMemory(int(self.wg_red * 8))
                k1 = stage1(
                    self.queue,
                    (int(self.wg_red**2),),
                    (int(self.wg_red),),
                    input_array.data,
                    self.cl_mem["tmp_max_min_d"].data,
                    size,
                    local_mem,
                )
                k2 = stage2(
                    self.queue,
                    (int(self.wg_red),),
                    (int(self.wg_red),),
                    self.cl_mem["tmp_max_min_d"].data,
                    self.cl_mem["max_min_d"].data,
                    local_mem,
                )

                events += [
                    EventDescription("max_min_stage1", k1),
                    EventDescription("max_min_stage2", k2),
                ]

            evt = self.kernels.normalize_image(
                self.queue,
                (self.shape[1], self.shape[0]),
                None,
                input_array.data,
                output_array.data,
                numpy.int32(self.shape[1]),
                numpy.int32(self.shape[0]),
                self.cl_mem["max_min_d"].data,
                numpy.float32(mini),
                numpy.float32(maxi),
            )
            events.append(EventDescription("normalize", evt))
        if self.profile:
            self.events += events

        if out is None:
            res = output_array.get()
            return res
        else:
            output_array.finish()
            return output_array
//  [python]
def normalize(self, img, mini=0.0, maxi=1.0, copy=True, out=None):
        """Scale the intensity of the image so that the minimum is 0 and the
        maximum is 1.0 (or any value suggested).

        :param img: numpy array or pyopencl array of dim 2 or 3 and of type float
        :param mini: Expected minimum value
        :param maxi: expected maxiumum value
        :param copy: set to False to use directly the input buffer
        :param out: provides an output buffer. prevents a copy D->H

        This uses a min/max reduction in two stages plus a map operation
        """
        assert img.shape == self.buffer_shape
        events = []
        with self.sem:
            input_array, output_array = self._get_in_out_buffers(img, copy, out)
            size = numpy.int32(numpy.prod(self.shape))
            if self.wg_red == 1:
                #  Probably on MacOS CPU WG==1 --> serial code.
                kernel = self.kernels.get_kernel("max_min_serial")
                evt = kernel(
                    self.queue,
                    (1,),
                    (1,),
                    input_array.data,
                    size,
                    self.cl_mem["max_min_d"].data,
                )
                ed = EventDescription("max_min_serial", evt)
                events.append(ed)
            else:
                stage1 = self.kernels.max_min_reduction_stage1
                stage2 = self.kernels.max_min_reduction_stage2
                local_mem = pyopencl.LocalMemory(int(self.wg_red * 8))
                k1 = stage1(
                    self.queue,
                    (int(self.wg_red**2),),
                    (int(self.wg_red),),
                    input_array.data,
                    self.cl_mem["tmp_max_min_d"].data,
                    size,
                    local_mem,
                )
                k2 = stage2(
                    self.queue,
                    (int(self.wg_red),),
                    (int(self.wg_red),),
                    self.cl_mem["tmp_max_min_d"].data,
                    self.cl_mem["max_min_d"].data,
                    local_mem,
                )

                events += [
                    EventDescription("max_min_stage1", k1),
                    EventDescription("max_min_stage2", k2),
                ]

            evt = self.kernels.normalize_image(
                self.queue,
                (self.shape[1], self.shape[0]),
                None,
                input_array.data,
                output_array.data,
                numpy.int32(self.shape[1]),
                numpy.int32(self.shape[0]),
                self.cl_mem["max_min_d"].data,
                numpy.float32(mini),
                numpy.float32(maxi),
            )
            events.append(EventDescription("normalize", evt))
        if self.profile:
            self.events += events

        if out is None:
            res = output_array.get()
            return res
        else:
            output_array.finish()
            return output_array
//  [python]
def __init__(
        self,
        image,
        patch_size: int = 32,
        pixel_masking_probability: float = 0.3,
    ):
        """

        Parameters
        ----------
        image
        """

        self.image = torch.tensor(image)
        self.patch_size = patch_size

        # print("shape: ", image.shape)

        self.patch_slicing_objects = random_patches(
            image=image,
            patch_size=patch_size,
            nb_patches_per_image=min(image.size / n.prod(patch_size), 1040),
        )
        self.p = pixel_masking_probability
//  [python]
def __init__(
        self,
        image,
        patch_size: int = 32,
        pixel_masking_probability: float = 0.3,
    ):
        """

        Parameters
        ----------
        image
        """

        self.image = torch.tensor(image)
        self.patch_size = patch_size

        # print("shape: ", image.shape)

        self.patch_slicing_objects = random_patches(
            image=image,
            patch_size=patch_size,
            nb_patches_per_image=min(image.size / numpy.prod(patch_size), 1040),
        )
        self.p = pixel_masking_probability
//  [python]
def _update_pointers(self):
        # TODO: call this anytime the _grid_id changes and anytime the structure_change_count change
        self._isalive()
        if self._species is None:
            return
        self._seg_indices = self._locate_segments()
        from .geometry import _surface_areas1d

        grid_list = 0
        grid_indices = []
        neuron_pointers = []
        stateo = "_ref_" + self._species + "o"
        for sec, indices in self._seg_indices.items():
            for seg, i in zip(sec, indices):
                if i is not None:
                    grid_indices.append(i)
                    neuron_pointers.append(getattr(seg, stateo))
        _set_grid_concentrations(
            grid_list, self._grid_id, grid_indices, neuron_pointers
        )
        if self._charge == 0 or isinstance(
            _defined_species[self._species][self._region](), Parameter
        ):
            _set_grid_currents(grid_list, self._grid_id, [], [], [])
        else:
            tenthousand_over_charge_faraday = 10000.0 / (self._charge * h.FARADAY)
            scale_factor = tenthousand_over_charge_faraday / (n.prod(self._dx))
            ispecies = "_ref_i" + self._species
            neuron_pointers = []
            scale_factors = []
            for sec, indices in self._seg_indices.items():
                for seg, surface_area, i in zip(sec, _surface_areas1d(sec), indices):
                    if i is not None:
                        neuron_pointers.append(getattr(seg, ispecies))
                        scale_factors.append(float(scale_factor * surface_area))
            # TODO: MultiCompartment reactions ?
            _set_grid_currents(
                grid_list, self._grid_id, grid_indices, neuron_pointers, scale_factors
            )
//  [python]
def _update_pointers(self):
        # TODO: call this anytime the _grid_id changes and anytime the structure_change_count change
        self._isalive()
        if self._species is None:
            return
        self._seg_indices = self._locate_segments()
        from .geometry import _surface_areas1d

        grid_list = 0
        grid_indices = []
        neuron_pointers = []
        stateo = "_ref_" + self._species + "o"
        for sec, indices in self._seg_indices.items():
            for seg, i in zip(sec, indices):
                if i is not None:
                    grid_indices.append(i)
                    neuron_pointers.append(getattr(seg, stateo))
        _set_grid_concentrations(
            grid_list, self._grid_id, grid_indices, neuron_pointers
        )
        if self._charge == 0 or isinstance(
            _defined_species[self._species][self._region](), Parameter
        ):
            _set_grid_currents(grid_list, self._grid_id, [], [], [])
        else:
            tenthousand_over_charge_faraday = 10000.0 / (self._charge * h.FARADAY)
            scale_factor = tenthousand_over_charge_faraday / (numpy.prod(self._dx))
            ispecies = "_ref_i" + self._species
            neuron_pointers = []
            scale_factors = []
            for sec, indices in self._seg_indices.items():
                for seg, surface_area, i in zip(sec, _surface_areas1d(sec), indices):
                    if i is not None:
                        neuron_pointers.append(getattr(seg, ispecies))
                        scale_factors.append(float(scale_factor * surface_area))
            # TODO: MultiCompartment reactions ?
            _set_grid_currents(
                grid_list, self._grid_id, grid_indices, neuron_pointers, scale_factors
            )
//  [python]
def __init__(
            self,
            parent: DataSegment,
            subset_definition: Tuple[slice, ...],
            coordinate_basis: str,
            squeeze: bool = True,
            close_parent: bool = True):
        """
        Parameters
        ----------
        parent : DataSegment
        subset_definition : Tuple[slice, ...]
        coordinate_basis : str
            The coordinate basis for the subset definition, it should be one of
            `('raw', 'formatted')`.
        squeeze: bool
            Eliminate the dimensions that are size 1 in the subset?
        close_parent : bool
            Call parent.close() when close is called?
        """

        self._close_parent = None
        self.close_parent = close_parent
        self._original_formatted_indices = None  # the original indices matched to the new, with entry -1 when flat
        self._original_raw_indices = None
        self._formatted_subset_definition = None
        self._raw_subset_definition = None
        self._squeeze = squeeze
        self._parent = parent

        raw_shape, formatted_shape = self._validate_subset_definition(
            subset_definition, coordinate_basis)
        DataSegment.__init__(
            self, parent.raw_dtype, raw_shape, parent.formatted_dtype, formatted_shape,
            mode=parent.mode)
        self._pixels_written = 0
        if self.mode == 'w':
            self._expected_pixels_written = int(numpy.prod(raw_shape))
        else:
            self._expected_pixels_written = 0
//  [python]
def get_optimal_chunks(shape, dtype, 
        target_chunk_bytes=2e6, 
        max_chunk_length=None):
    """
    Common shapes
    T,D
    T,N,D
    T,H,W,C
    T,N,H,W,C
    """
    itemsize = np.dtype(dtype).itemsize
    # reversed
    rshape = list(shape[::-1])
    if max_chunk_length is not None:
        rshape[-1] = int(max_chunk_length)
    split_idx = len(shape)-1
    for i in range(len(shape)-1):
        this_chunk_bytes = itemsize * reduce(lambda x, y: x*y, rshape[:i])
        next_chunk_bytes = itemsize * reduce(lambda x, y: x*y, rshape[:i+1])
        if this_chunk_bytes <= target_chunk_bytes \
            and next_chunk_bytes > target_chunk_bytes:
            split_idx = i

    rchunks = rshape[:split_idx]
    item_chunk_bytes = itemsize * reduce(lambda x, y: x*y, rshape[:split_idx])
    this_max_chunk_length = rshape[split_idx]
    next_chunk_length = min(this_max_chunk_length, math.ceil(
            target_chunk_bytes / item_chunk_bytes))
    rchunks.append(next_chunk_length)
    len_diff = len(shape) - len(rchunks)
    rchunks.extend([1] * len_diff)
    chunks = tuple(rchunks[::-1])
    # print(np.prod(chunks) * itemsize / target_chunk_bytes)
    return chunks
//  [python]
def get_optimal_chunks(shape, dtype, 
        target_chunk_bytes=2e6, 
        max_chunk_length=None):
    """
    Common shapes
    T,D
    T,N,D
    T,H,W,C
    T,N,H,W,C
    """
    itemsize = np.dtype(dtype).itemsize
    # reversed
    rshape = list(shape[::-1])
    if max_chunk_length is not None:
        rshape[-1] = int(max_chunk_length)
    split_idx = len(shape)-1
    for i in range(len(shape)-1):
        this_chunk_bytes = itemsize * np.prod(rshape[:i])
        next_chunk_bytes = itemsize * np.prod(rshape[:i+1])
        if this_chunk_bytes <= target_chunk_bytes \
            and next_chunk_bytes > target_chunk_bytes:
            split_idx = i

    rchunks = rshape[:split_idx]
    item_chunk_bytes = itemsize * np.prod(rshape[:split_idx])
    this_max_chunk_length = rshape[split_idx]
    next_chunk_length = min(this_max_chunk_length, math.ceil(
            target_chunk_bytes / item_chunk_bytes))
    rchunks.append(next_chunk_length)
    len_diff = len(shape) - len(rchunks)
    rchunks.extend([1] * len_diff)
    chunks = tuple(rchunks[::-1])
    # print(np.prod(chunks) * itemsize / target_chunk_bytes)
    return chunks
//  [python]
def uniform(cls, BoxSize, comm=MPI.COMM_WORLD, periodic=True):
        ndim = len(BoxSize)

        # compute a optimal shape where each domain is as cubical as possible

        r = (1.0 * comm.size / n.prod(BoxSize) * min(BoxSize)) ** (1.0 / ndim)
        shape = [ r * (BoxSize[i] / min(BoxSize)) for i in range(ndim)]
        shape = numpy.array(shape)
        imax = shape.argmax()
        shape = numpy.int32(shape)
        shape[shape < 1] = 1
        shape[imax] = 1
        shape[imax] = comm.size // n.prod(shape)
        assert n.prod(shape) <= comm.size

        edges = []
        for i in range(ndim):
            edges.append(numpy.linspace(0, BoxSize[i], shape[i] + 1, endpoint=True))
        return cls(edges, comm, periodic)
//  [python]
def uniform(cls, BoxSize, comm=MPI.COMM_WORLD, periodic=True):
        ndim = len(BoxSize)

        # compute a optimal shape where each domain is as cubical as possible

        r = (1.0 * comm.size / numpy.prod(BoxSize) * min(BoxSize)) ** (1.0 / ndim)
        shape = [ r * (BoxSize[i] / min(BoxSize)) for i in range(ndim)]
        shape = numpy.array(shape)
        imax = shape.argmax()
        shape = numpy.int32(shape)
        shape[shape < 1] = 1
        shape[imax] = 1
        shape[imax] = comm.size // numpy.prod(shape)
        assert numpy.prod(shape) <= comm.size

        edges = []
        for i in range(ndim):
            edges.append(numpy.linspace(0, BoxSize[i], shape[i] + 1, endpoint=True))
        return cls(edges, comm, periodic)
//  [python]
def _loglikelihood(self):
        """Returns the log pdf of the eggbox function.
        """
        return (2 + n.prod(n.cos([
            self.current_params[p]/2. for p in self.variable_params]))) ** 5
//  [python]
def _loglikelihood(self):
        """Returns the log pdf of the eggbox function.
        """
        return (2 + numpy.prod(numpy.cos([
            self.current_params[p]/2. for p in self.variable_params]))) ** 5
//  [python]
def read_raw(
            self,
            subscript: Union[None, int, slice, Sequence[Union[int, slice, Tuple[int, ...]]]],
            squeeze=True) -> numpy.ndarray:
        self._validate_closed()
        subscript, out_shape = get_subscript_result_size(subscript, self.raw_shape)

        init_slice = subscript[0]
        init_reverse = (init_slice.step < 0)
        if init_reverse:
            init_slice = _reverse_slice(init_slice)

        pixel_per_row = 1 if self.formatted_ndim == 1 else int(numpy.multiply.reduce(self.raw_shape[1:]))
        row_stride = self.raw_dtype.itemsize*pixel_per_row

        start_row = init_slice.start
        rows = init_slice.stop - init_slice.start

        # read the whole contiguous chunk from start_row up to the final row
        # seek to the proper start location
        start_loc = self._data_offset + start_row*row_stride
        self.file_object.seek(start_loc, os.SEEK_SET)
        total_size = rows*row_stride
        # read our data
        data = self.file_object.read(total_size)
        if len(data) != total_size:
            raise ValueError(
                'Tried to read {} bytes of data, but received {}.\n'
                'The most likely reason for this is a malformed chipper, \n'
                'which attempts to read more data than the file contains'.format(total_size, len(data)))
        # define temp array from this data
        data = numpy.frombuffer(data, self._raw_dtype, rows*pixel_per_row)
        data = numpy.reshape(data, (rows, ) + self.raw_shape[1:])
        # extract our data
        out = data[(slice(None, None, init_slice.step), ) + subscript[1:]]
        out = numpy.reshape(out, out_shape)
        if init_reverse:
            out = numpy.flip(out, axis=0)

        if squeeze:
            out = numpy.copy(numpy.squeeze(out))
        else:
            out = numpy.copy(out)
        del data
        return out
//  [python]
def read_raw(
            self,
            subscript: Union[None, int, slice, Sequence[Union[int, slice, Tuple[int, ...]]]],
            squeeze=True) -> numpy.ndarray:
        self._validate_closed()
        subscript, out_shape = get_subscript_result_size(subscript, self.raw_shape)

        init_slice = subscript[0]
        init_reverse = (init_slice.step < 0)
        if init_reverse:
            init_slice = _reverse_slice(init_slice)

        pixel_per_row = 1 if self.formatted_ndim == 1 else int(numpy.prod(self.raw_shape[1:]))
        row_stride = self.raw_dtype.itemsize*pixel_per_row

        start_row = init_slice.start
        rows = init_slice.stop - init_slice.start

        # read the whole contiguous chunk from start_row up to the final row
        # seek to the proper start location
        start_loc = self._data_offset + start_row*row_stride
        self.file_object.seek(start_loc, os.SEEK_SET)
        total_size = rows*row_stride
        # read our data
        data = self.file_object.read(total_size)
        if len(data) != total_size:
            raise ValueError(
                'Tried to read {} bytes of data, but received {}.\n'
                'The most likely reason for this is a malformed chipper, \n'
                'which attempts to read more data than the file contains'.format(total_size, len(data)))
        # define temp array from this data
        data = numpy.frombuffer(data, self._raw_dtype, rows*pixel_per_row)
        data = numpy.reshape(data, (rows, ) + self.raw_shape[1:])
        # extract our data
        out = data[(slice(None, None, init_slice.step), ) + subscript[1:]]
        out = numpy.reshape(out, out_shape)
        if init_reverse:
            out = numpy.flip(out, axis=0)

        if squeeze:
            out = numpy.copy(numpy.squeeze(out))
        else:
            out = numpy.copy(out)
        del data
        return out
//  [python]
def get_required_memory_size(self, view_model):
        # type: (FFTAdapterModel) -> int
        """
        Returns the required memory to be able to run the adapter.
        """
        input_size = numpy.prod(self.input_shape) * 8.0
        output_size = self.result_size(self.input_shape, view_model.segment_length,
                                       self.input_time_series_index.sample_period)
        total_free_memory = psutil.virtual_memory().free + psutil.swap_memory().free
        total_required_memory = input_size + output_size
        while total_required_memory / self.memory_factor / total_free_memory > 0.8:
            self.memory_factor += 1
        return total_required_memory / self.memory_factor
//  [python]
def result_size(self, input_shape, segment_length, sample_period):
        """
        Returns the storage size in Bytes of the main result (complex array) of
        the FFT.
        """
        result_size = numpy.prod(self.result_shape(input_shape, segment_length,
                                                   sample_period)) * 2.0 * 8.0  # complex*Bytes
        return result_size
//  [python]
def test_sparse_nd_storage_fallback():
    def check_output_fallback(shape):
        ones = mx.nd.ones(shape)
        out = mx.nd.zeros(shape=shape, stype='csr')
        mx.nd.broadcast_add(ones, ones * 2, out=out)
        assert(np.sum(out.asnumpy() - 3) == 0)

    def check_input_fallback(shape):
        ones = mx.nd.ones(shape)
        out = mx.nd.broadcast_add(ones.tostype('csr'), ones.tostype('row_sparse'))
        assert(np.sum(out.asnumpy() - 2) == 0)

    def check_fallback_with_temp_resource(shape):
        ones = mx.nd.ones(shape)
        out = mx.nd.sum(ones)
        assert(out.asscalar() == np.prod(shape))

    shape = rand_shape_2d()
    check_output_fallback(shape)
    check_input_fallback(shape)
    check_fallback_with_temp_resource(shape)
//  [python]
def bytes_to_arrays(serialized_arr: bytearray) -> List[numpy.array]:
    """
    :param serialized_arr: bytearray representation of list of numpy arrays
    :return: List of numpy arrays decoded from input
    """
    sep = "|".encode("utf-8")
    arrays = []
    i_start = 0
    while i_start < len(serialized_arr) - 1:
        i_0 = serialized_arr.find(sep, i_start)
        i_1 = serialized_arr.find(sep, i_0 + 1)
        arr_dtype = numpy.dtype(serialized_arr[i_start:i_0].decode("utf-8"))
        arr_shape = tuple(
            [int(a) for a in serialized_arr[i_0 + 1 : i_1].decode("utf-8").split(",")]
        )
        arr_num_bytes = reduce(lambda x, y: x * y, arr_shape) * arr_dtype.itemsize
        arr_str = serialized_arr[i_1 + 1 : arr_num_bytes + (i_1 + 1)]
        arr = numpy.frombuffer(arr_str, dtype=arr_dtype).reshape(arr_shape)
        arrays.append(arr.copy())

        i_start = i_1 + arr_num_bytes + 1
    return arrays
//  [python]
def bytes_to_arrays(serialized_arr: bytearray) -> List[numpy.array]:
    """
    :param serialized_arr: bytearray representation of list of numpy arrays
    :return: List of numpy arrays decoded from input
    """
    sep = "|".encode("utf-8")
    arrays = []
    i_start = 0
    while i_start < len(serialized_arr) - 1:
        i_0 = serialized_arr.find(sep, i_start)
        i_1 = serialized_arr.find(sep, i_0 + 1)
        arr_dtype = numpy.dtype(serialized_arr[i_start:i_0].decode("utf-8"))
        arr_shape = tuple(
            [int(a) for a in serialized_arr[i_0 + 1 : i_1].decode("utf-8").split(",")]
        )
        arr_num_bytes = numpy.prod(arr_shape) * arr_dtype.itemsize
        arr_str = serialized_arr[i_1 + 1 : arr_num_bytes + (i_1 + 1)]
        arr = numpy.frombuffer(arr_str, dtype=arr_dtype).reshape(arr_shape)
        arrays.append(arr.copy())

        i_start = i_1 + arr_num_bytes + 1
    return arrays
//  [python]
def bytecount(self, dtype=float):
        """Get the number of bytes used by a grid of specified dtype."""
        return long(reduce(lambda x,y:x*y, self.n_c)) * np.array(1, dtype).itemsize
//  [python]
def bytecount(self, dtype=float):
        """Get the number of bytes used by a grid of specified dtype."""
        return long(np.prod(self.n_c)) * np.array(1, dtype).itemsize
//  [python]
def wannier_matrix(self, psit_nG, psit_nG1, c, G, nbands=None):
        """Wannier localization integrals

        The soft part of Z is given by (Eq. 27 ref1)::

            ~       ~     -i G.r   ~
            Z   = <psi | e      |psi >
             nm       n             m
                    
        G is 1/N_c (plus 1 if k-points distances should be wrapped over
        the Brillouin zone), where N_c is the number of k-points along
        axis c, psit_nG and psit_nG1 are the set of wave functions for
        the two different spin/kpoints in question.

        ref1: Thygesen et al, Phys. Rev. B 72, 125119 (2005) 
        """
        same_wave = False
        if psit_nG is psit_nG1:
            same_wave = True

        if nbands is None:
            nbands = len(psit_nG)
        
        def get_slice(c, g, psit_nG):
            if c == 0:
                slice_nG = psit_nG[:nbands, g].copy()
            elif c == 1:
                slice_nG = psit_nG[:nbands, :, g].copy()
            else:
                slice_nG = psit_nG[:nbands, :, :, g].copy()
            return slice_nG.reshape((nbands, reduce(lambda x, y: x*y, slice_nG.shape[1:])))
        
        Z_nn = np.zeros((nbands, nbands), complex)
        for g in range(self.n_c[c]):
            A_nG = get_slice(c, g, psit_nG)
                
            if same_wave:
                B_nG = A_nG
            else:
                B_nG = get_slice(c, g, psit_nG1)
                
            e = exp(-2.j * pi * G * (g + self.beg_c[c]) / self.N_c[c])
            Z_nn += e * np.dot(A_nG.conj(), B_nG.T) * self.dv
            
        return Z_nn
//  [python]
def wannier_matrix(self, psit_nG, psit_nG1, c, G, nbands=None):
        """Wannier localization integrals

        The soft part of Z is given by (Eq. 27 ref1)::

            ~       ~     -i G.r   ~
            Z   = <psi | e      |psi >
             nm       n             m
                    
        G is 1/N_c (plus 1 if k-points distances should be wrapped over
        the Brillouin zone), where N_c is the number of k-points along
        axis c, psit_nG and psit_nG1 are the set of wave functions for
        the two different spin/kpoints in question.

        ref1: Thygesen et al, Phys. Rev. B 72, 125119 (2005) 
        """
        same_wave = False
        if psit_nG is psit_nG1:
            same_wave = True

        if nbands is None:
            nbands = len(psit_nG)
        
        def get_slice(c, g, psit_nG):
            if c == 0:
                slice_nG = psit_nG[:nbands, g].copy()
            elif c == 1:
                slice_nG = psit_nG[:nbands, :, g].copy()
            else:
                slice_nG = psit_nG[:nbands, :, :, g].copy()
            return slice_nG.reshape((nbands, np.prod(slice_nG.shape[1:])))
        
        Z_nn = np.zeros((nbands, nbands), complex)
        for g in range(self.n_c[c]):
            A_nG = get_slice(c, g, psit_nG)
                
            if same_wave:
                B_nG = A_nG
            else:
                B_nG = get_slice(c, g, psit_nG1)
                
            e = exp(-2.j * pi * G * (g + self.beg_c[c]) / self.N_c[c])
            Z_nn += e * np.dot(A_nG.conj(), B_nG.T) * self.dv
            
        return Z_nn
//  [python]
def test_concatenations_and_stacking(self):
        for backend in imp_op_backends:
            print('testing shapes for ', backend.framework_name)
            for n_arrays in [1, 2, 5]:
                shapes = [[], [1], [1, 1], [2, 3, 5, 7], [1] * 6]
                for shape in shapes:
                    if (backend.framework_name == 'jittor')\
                            and len(shape) == 0:
                        # jittor stores scalar in 1d array
                        continue
                    arrays1 = [numpy.arange(i, i + numpy.prod(shape)).reshape(shape) for i in range(n_arrays)]
                    arrays2 = [backend.from_numpy(array) for array in arrays1]
                    result0 = numpy.asarray(arrays1)
                    result1 = rearrange(arrays1, '...->...')
                    result2 = rearrange(arrays2, '...->...')
                    assert numpy.array_equal(result0, result1)
                    assert numpy.array_equal(result1, backend.to_numpy(result2))

                    result1 = rearrange(arrays1, 'b ... -> ... b')
                    result2 = rearrange(arrays2, 'b ... -> ... b')
                    assert numpy.array_equal(result1, backend.to_numpy(result2))
//  [python]
def test_enumerating_directions(self):
        for backend in imp_op_backends:
            print('testing directions for', backend.framework_name)
            for shape in [[], [1], [1, 1, 1], [2, 3, 5, 7]]:
                if (backend.framework_name == 'jittor')\
                        and len(shape) == 0:
                    # jittor stores scalar in 1d array
                    continue
                x = numpy.arange(numpy.prod(shape)).reshape(shape)
                axes1 = _enumerate_directions(x)
                axes2 = _enumerate_directions(backend.from_numpy(x))
                assert len(axes1) == len(axes2) == len(shape)
                for ax1, ax2 in zip(axes1, axes2):
                    ax2 = backend.to_numpy(ax2)
                    assert ax1.shape == ax2.shape
                    assert numpy.allclose(ax1, ax2)
//  [python]
def test_fft_kpq(nalpha):
    # Create regular grid.
    nmax = 1
    mesh = [2*nmax+1]*3
    grid, eigs = generate_fft_grid(mesh)

    qmax = 2*nmax
    qmesh = [2*qmax+1]*3
    qgrid, qeigs = generate_fft_grid(qmesh)
    # Create wavefunction
    nbasis = len(grid)

    numpy.random.seed(7)
    psi = get_random_wavefunction((nalpha,nalpha), nbasis)
    I = get_random_wavefunction((nalpha,nalpha), nbasis)
    trial = I[:,:nalpha].conj()

    # Select lowest energy states for trial
    G, Gh = gab_mod(trial,psi[:,:nalpha])
    nqgrid = numpy.multiply.reduce(qmesh)

    # # Check by direct convolution f(q) = \sum_G Psi[G+Q] Gh[G].
    fq_direct = numpy.zeros(nqgrid, dtype=numpy.complex128)
    for i, q in enumerate(qgrid):
        for j, k in enumerate(grid):
            kpq = k + q
            ikpq = lookup(kpq, grid)
            if ikpq is not None:
                fq_direct[i] += trial[ikpq,0] * Gh[0,j]

    trial_grid = numpy.flip(trial[:,0]).reshape(mesh)
    Gh_grid = Gh[0,:].reshape(mesh)

    # Check by fft convolve
    # Compare to fq
    fq_conv = numpy.zeros(nqgrid, dtype=numpy.complex128)
    fq_conv += nqgrid*convolve(Gh_grid, trial_grid, mesh)
    fq_conv = numpy.flip(fq_conv)

    import scipy.signal
    fq_conv_sc = numpy.flip(scipy.signal.fftconvolve(Gh_grid,trial_grid).ravel())

    import matplotlib.pyplot as pl
    pl.plot(fq_conv, label='fft')
    pl.plot(fq_conv_sc, label='fft_scipy')
    pl.plot(fq_direct, label='direct')
    # pl.plot(fq, label='from_gf')
    pl.legend()
    pl.show()
//  [python]
def test_fft_kpq(nalpha):
    # Create regular grid.
    nmax = 1
    mesh = [2*nmax+1]*3
    grid, eigs = generate_fft_grid(mesh)

    qmax = 2*nmax
    qmesh = [2*qmax+1]*3
    qgrid, qeigs = generate_fft_grid(qmesh)
    # Create wavefunction
    nbasis = len(grid)

    numpy.random.seed(7)
    psi = get_random_wavefunction((nalpha,nalpha), nbasis)
    I = get_random_wavefunction((nalpha,nalpha), nbasis)
    trial = I[:,:nalpha].conj()

    # Select lowest energy states for trial
    G, Gh = gab_mod(trial,psi[:,:nalpha])
    nqgrid = numpy.prod(qmesh)

    # # Check by direct convolution f(q) = \sum_G Psi[G+Q] Gh[G].
    fq_direct = numpy.zeros(nqgrid, dtype=numpy.complex128)
    for i, q in enumerate(qgrid):
        for j, k in enumerate(grid):
            kpq = k + q
            ikpq = lookup(kpq, grid)
            if ikpq is not None:
                fq_direct[i] += trial[ikpq,0] * Gh[0,j]

    trial_grid = numpy.flip(trial[:,0]).reshape(mesh)
    Gh_grid = Gh[0,:].reshape(mesh)

    # Check by fft convolve
    # Compare to fq
    fq_conv = numpy.zeros(nqgrid, dtype=numpy.complex128)
    fq_conv += nqgrid*convolve(Gh_grid, trial_grid, mesh)
    fq_conv = numpy.flip(fq_conv)

    import scipy.signal
    fq_conv_sc = numpy.flip(scipy.signal.fftconvolve(Gh_grid,trial_grid).ravel())

    import matplotlib.pyplot as pl
    pl.plot(fq_conv, label='fft')
    pl.plot(fq_conv_sc, label='fft_scipy')
    pl.plot(fq_direct, label='direct')
    # pl.plot(fq, label='from_gf')
    pl.legend()
    pl.show()
//  [python]
def test_fft_kmq(nalpha):
    # Create regular grid.
    nmax = 1
    mesh = [2*nmax+1]*3
    grid, eigs = generate_fft_grid(mesh)
    qmax = 2*nmax
    qmesh = [2*qmax+1]*3
    qgrid, qeigs = generate_fft_grid(qmesh)
    # Create wavefunction
    nbasis = len(grid)
    # numpy.random.seed(7)
    psi = get_random_wavefunction((nalpha,nalpha), nbasis)
    I = get_random_wavefunction((nalpha,nalpha), nbasis)

    # Select lowest energy states for trial
    trial = I[:,:nalpha].conj()
    G, Gh = gab_mod(trial,psi[:,:nalpha])
    nqgrid = n.prod(qmesh)

    # # Check by direct convolution f(q) = \sum_G Psi[G-Q] Gh[G].
    fq_direct = numpy.zeros(nqgrid, dtype=numpy.complex128)
    for iq, q in enumerate(qgrid):
        for i, g in enumerate(grid):
            gmq = g - q
            igmq = lookup(gmq, grid)
            if igmq is not None:
                fq_direct[iq] += trial[igmq,0] * Gh[0,i]

    trial_grid = trial[:,0].reshape(mesh)
    Gh_grid = numpy.flip(Gh[0,:]).reshape(mesh)

    # Check by fft convolve
    # Compare to fq
    fq_conv = numpy.zeros(nqgrid, dtype=numpy.complex128)
    fq_conv += nqgrid*convolve(trial_grid, Gh_grid, mesh)
    fq_conv = numpy.flip(fq_conv)

    import scipy.signal
    fq_conv_sc = numpy.flip(scipy.signal.fftconvolve(trial_grid, Gh_grid).ravel())

    import matplotlib.pyplot as pl
    pl.plot(fq_conv, label='fft')
    pl.plot(fq_conv_sc, label='fft_scipy')
    pl.plot(fq_direct, label='direct')
    pl.legend()
    pl.show()
//  [python]
def test_fft_kmq(nalpha):
    # Create regular grid.
    nmax = 1
    mesh = [2*nmax+1]*3
    grid, eigs = generate_fft_grid(mesh)
    qmax = 2*nmax
    qmesh = [2*qmax+1]*3
    qgrid, qeigs = generate_fft_grid(qmesh)
    # Create wavefunction
    nbasis = len(grid)
    # numpy.random.seed(7)
    psi = get_random_wavefunction((nalpha,nalpha), nbasis)
    I = get_random_wavefunction((nalpha,nalpha), nbasis)

    # Select lowest energy states for trial
    trial = I[:,:nalpha].conj()
    G, Gh = gab_mod(trial,psi[:,:nalpha])
    nqgrid = numpy.prod(qmesh)

    # # Check by direct convolution f(q) = \sum_G Psi[G-Q] Gh[G].
    fq_direct = numpy.zeros(nqgrid, dtype=numpy.complex128)
    for iq, q in enumerate(qgrid):
        for i, g in enumerate(grid):
            gmq = g - q
            igmq = lookup(gmq, grid)
            if igmq is not None:
                fq_direct[iq] += trial[igmq,0] * Gh[0,i]

    trial_grid = trial[:,0].reshape(mesh)
    Gh_grid = numpy.flip(Gh[0,:]).reshape(mesh)

    # Check by fft convolve
    # Compare to fq
    fq_conv = numpy.zeros(nqgrid, dtype=numpy.complex128)
    fq_conv += nqgrid*convolve(trial_grid, Gh_grid, mesh)
    fq_conv = numpy.flip(fq_conv)

    import scipy.signal
    fq_conv_sc = numpy.flip(scipy.signal.fftconvolve(trial_grid, Gh_grid).ravel())

    import matplotlib.pyplot as pl
    pl.plot(fq_conv, label='fft')
    pl.plot(fq_conv_sc, label='fft_scipy')
    pl.plot(fq_direct, label='direct')
    pl.legend()
    pl.show()
//  [python]
def size(self) -> int:
        """
        Return an int representing the number of elements in this object.

        Return the number of rows if Series. Otherwise return the number of
        rows times number of columns if DataFrame.

        See Also
        --------
        ndarray.size : Number of elements in the array.

        Examples
        --------
        >>> s = pd.Series({"a": 1, "b": 2, "c": 3})
        >>> s.size
        3

        >>> df = pd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
        >>> df.size
        4
        """

        return int(np.multiply.reduce(self.shape))
//  [python]
def size(self) -> int:
        """
        Return an int representing the number of elements in this object.

        Return the number of rows if Series. Otherwise return the number of
        rows times number of columns if DataFrame.

        See Also
        --------
        ndarray.size : Number of elements in the array.

        Examples
        --------
        >>> s = pd.Series({"a": 1, "b": 2, "c": 3})
        >>> s.size
        3

        >>> df = pd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
        >>> df.size
        4
        """

        return int(np.prod(self.shape))
//  [python]
def test_reduction_stress_imperatives(self):
        for backend in imp_op_backends:
            print('Stress-testing reduction for ', backend.framework_name)
            for reduction in _reductions + ('rearrange',):
                dtype = 'int64'
                coincide = numpy.array_equal
                if reduction in ['mean', 'prod']:
                    dtype = 'float64'
                    coincide = numpy.allclose
                for n_axes in range(11):
                    shape = numpy.random.randint(2, 4, size=n_axes)
                    permutation = numpy.random.permutation(n_axes)
                    skipped = 0 if reduction == 'rearrange' else numpy.random.randint(n_axes + 1)
                    left = ' '.join('x' + str(i) for i in range(n_axes))
                    right = ' '.join('x' + str(i) for i in permutation[skipped:])
                    pattern = left + '->' + right
                    x = numpy.arange(1, 1 + numpy.prod(shape), dtype=dtype).reshape(shape)
                    if reduction == 'prod':
                        x /= x.mean()  # to avoid overflows
                    result1 = reduce(x, pattern, reduction=reduction)
                    result2 = x.transpose(permutation)
                    if skipped > 0:
                        result2 = getattr(result2, reduction)(axis=tuple(range(skipped)))
                    assert coincide(result1, result2)
                    check_op_against_numpy(backend, x, pattern, reduction=reduction, axes_lengths={}, is_symbolic=False)
//  [python]
def test_rearrange_imperative(self):
        for backend in collect_test_backends(symbolic=False, layers=True):
            print('Test layer for ', backend.framework_name)

            for pattern, axes_lengths, input_shape, wrong_shapes in rearrangement_patterns:
                x = numpy.arange(n.prod(input_shape), dtype='float32').reshape(input_shape)
                result_numpy = rearrange(x, pattern, **axes_lengths)
                layer = backend.layers().Rearrange(pattern, **axes_lengths)
                for shape in wrong_shapes:
                    try:
                        layer(backend.from_numpy(numpy.zeros(shape, dtype='float32')))
                    except:
                        pass
                    else:
                        raise AssertionError('Failure expected')

                # simple pickling / unpickling
                layer2 = pickle.loads(pickle.dumps(layer))
                result1 = backend.to_numpy(layer(backend.from_numpy(x)))
                result2 = backend.to_numpy(layer2(backend.from_numpy(x)))
                assert numpy.allclose(result_numpy, result1)
                assert numpy.allclose(result1, result2)

                just_sum = backend.layers().Reduce('...->', reduction='sum')

               
                variable = backend.from_numpy(x)
                result = just_sum(layer(variable))

                if 'jittor' in backend.framework_name:
                    grad = backend.jittor.grad(result, variable)
                else:
                    result.backward()
                    grad = variable.grad

                assert numpy.allclose(backend.to_numpy(grad), 1)
//  [python]
def test_rearrange_imperative(self):
        for backend in collect_test_backends(symbolic=False, layers=True):
            print('Test layer for ', backend.framework_name)

            for pattern, axes_lengths, input_shape, wrong_shapes in rearrangement_patterns:
                x = numpy.arange(numpy.prod(input_shape), dtype='float32').reshape(input_shape)
                result_numpy = rearrange(x, pattern, **axes_lengths)
                layer = backend.layers().Rearrange(pattern, **axes_lengths)
                for shape in wrong_shapes:
                    try:
                        layer(backend.from_numpy(numpy.zeros(shape, dtype='float32')))
                    except:
                        pass
                    else:
                        raise AssertionError('Failure expected')

                # simple pickling / unpickling
                layer2 = pickle.loads(pickle.dumps(layer))
                result1 = backend.to_numpy(layer(backend.from_numpy(x)))
                result2 = backend.to_numpy(layer2(backend.from_numpy(x)))
                assert numpy.allclose(result_numpy, result1)
                assert numpy.allclose(result1, result2)

                just_sum = backend.layers().Reduce('...->', reduction='sum')

               
                variable = backend.from_numpy(x)
                result = just_sum(layer(variable))

                if 'jittor' in backend.framework_name:
                    grad = backend.jittor.grad(result, variable)
                else:
                    result.backward()
                    grad = variable.grad

                assert numpy.allclose(backend.to_numpy(grad), 1)