Python zarr.open_group方法代码示例

# 需要导入模块: import zarr [as 别名]
# 或者: from zarr import open_group [as 别名]
def initialize_output_zarr(out_dir, sad_stats, snps, target_ids, target_labels):
  """Initialize an output Zarr file for SAD stats."""

  num_targets = len(target_ids)
  num_snps = len(snps)

  sad_out = zarr.open_group('%s/sad.zarr' % out_dir, 'w')

  # write SNPs
  sad_out.create_dataset('snp', data=[snp.rsid for snp in snps], chunks=(32768,))

  # write targets
  sad_out.create_dataset('target_ids', data=target_ids, compressor=None)
  sad_out.create_dataset('target_labels', data=target_labels, compressor=None)

  # initialize SAD stats
  for sad_stat in sad_stats:
    sad_out.create_dataset(sad_stat,
        shape=(num_snps, num_targets),
        chunks=(128, num_targets),
        dtype='float16')

  return sad_out 

# 需要导入模块: import zarr [as 别名]
# 或者: from zarr import open_group [as 别名]
def test_zarr_multiscale():
    multiscale = [
        np.random.random((20, 20)),
        np.random.random((10, 10)),
        np.random.random((5, 5)),
    ]
    with TemporaryDirectory(suffix='.zarr') as fout:
        root = zarr.open_group(fout, 'a')
        for i in range(len(multiscale)):
            shape = 20 // 2 ** i
            z = root.create_dataset(str(i), shape=(shape,) * 2)
            z[:] = multiscale[i]
        multiscale_in = io.magic_imread([fout])
        assert len(multiscale) == len(multiscale_in)
        # Note: due to lazy loading, the next line needs to happen within
        # the context manager. Alternatively, we could convert to NumPy here.
        for images, images_in in zip(multiscale, multiscale_in):
            np.testing.assert_array_equal(images, images_in) 

# 需要导入模块: import zarr [as 别名]
# 或者: from zarr import open_group [as 别名]
def test_zarr_multiscale():
    viewer = ViewerModel()
    multiscale = [
        np.random.random((20, 20)),
        np.random.random((10, 10)),
        np.random.random((5, 5)),
    ]
    with TemporaryDirectory(suffix='.zarr') as fout:
        root = zarr.open_group(fout, 'a')
        for i in range(len(multiscale)):
            shape = 20 // 2 ** i
            z = root.create_dataset(str(i), shape=(shape,) * 2)
            z[:] = multiscale[i]
        viewer.open(fout, multiscale=True, plugin='builtins')
        assert len(viewer.layers) == 1
        assert len(multiscale) == len(viewer.layers[0].data)
        # Note: due to lazy loading, the next line needs to happen within
        # the context manager. Alternatively, we could convert to NumPy here.
        for images, images_in in zip(multiscale, viewer.layers[0].data):
            np.testing.assert_array_equal(images, images_in) 

# 需要导入模块: import zarr [as 别名]
# 或者: from zarr import open_group [as 别名]
def test_resize_zarr_dataset(self):
        @xs.process
        class P:
            arr = xs.variable(dims="x", intent="out")

        model = xs.Model({"p": P})

        in_ds = xs.create_setup(
            model=model, clocks={"clock": [0, 1, 2]}, output_vars={"p__arr": "clock"},
        )

        store = ZarrSimulationStore(in_ds, model)

        for step, size in zip([0, 1, 2], [1, 3, 2]):
            model.state[("p", "arr")] = np.ones(size)
            store.write_output_vars(-1, step)

        ztest = zarr.open_group(store.zgroup.store, mode="r")

        expected = np.array(
            [[1.0, np.nan, np.nan], [1.0, 1.0, 1.0], [1.0, 1.0, np.nan]]
        )
        np.testing.assert_array_equal(ztest.p__arr, expected) 

# 需要导入模块: import zarr [as 别名]
# 或者: from zarr import open_group [as 别名]
def __init__(
        self,
        path=None,
        num_flush_threads=0,
        compressor=DEFAULT_COMPRESSOR,
        chunk_size=1024,
        max_file_size=None,
    ):
        self._mode = self.BUILD_MODE
        self._num_flush_threads = num_flush_threads
        self._chunk_size = max(1, chunk_size)
        self._metadata_codec = numcodecs.JSON()
        self._compressor = compressor
        self.data = zarr.group()
        self.path = path
        if path is not None:
            store = self._new_lmdb_store(max_file_size)
            self.data = zarr.open_group(store=store, mode="w")
        self.data.attrs[FORMAT_NAME_KEY] = self.FORMAT_NAME
        self.data.attrs[FORMAT_VERSION_KEY] = self.FORMAT_VERSION
        self.data.attrs["uuid"] = str(uuid.uuid4())

        chunks = self._chunk_size
        provenances_group = self.data.create_group("provenances")
        provenances_group.create_dataset(
            "timestamp",
            shape=(0,),
            chunks=chunks,
            compressor=self._compressor,
            dtype=object,
            object_codec=self._metadata_codec,
        )
        provenances_group.create_dataset(
            "record",
            shape=(0,),
            chunks=chunks,
            compressor=self._compressor,
            dtype=object,
            object_codec=self._metadata_codec,
        ) 

# 需要导入模块: import zarr [as 别名]
# 或者: from zarr import open_group [as 别名]
def collect_zarr(file_name, out_dir, num_procs):
  final_zarr_file = '%s/%s' % (out_dir, file_name)

  # seed w/ job0
  job_zarr_file = '%s/job0/%s' % (out_dir, file_name)
  shutil.copytree(job_zarr_file, final_zarr_file)

  # open final
  final_zarr_open = zarr.open_group(final_zarr_file)

  for pi in range(1, num_procs):
    # open job
    job_zarr_file = '%s/job%d/%s' % (out_dir, pi, file_name)
    job_zarr_open = zarr.open_group(job_zarr_file, 'r')

    # append to final
    for key in final_zarr_open.keys():
      if key in ['percentiles', 'target_ids', 'target_labels']:
        # once is enough
        pass

      elif key[-4:] == '_pct':
        # average
        u_k1 = np.array(final_zarr_open[key])
        x_k = np.array(job_zarr_open[key])
        final_zarr_open[key] = u_k1 + (x_k - u_k1) / (pi+1)

      else:
        # append
        final_zarr_open[key].append(job_zarr_open[key]) 

# 需要导入模块: import zarr [as 别名]
# 或者: from zarr import open_group [as 别名]
def update(self, output_path: str, global_attrs: Dict[str, Any] = None, **kwargs):
        if global_attrs:
            import zarr
            ds = zarr.open_group(output_path, mode='r+', **kwargs)
            ds.attrs.update(global_attrs)


# noinspection PyAbstractClass 

# 需要导入模块: import zarr [as 别名]
# 或者: from zarr import open_group [as 别名]
def append_time_slice(store: Union[str, MutableMapping],
                      time_slice: xr.Dataset,
                      chunk_sizes: Dict[str, int] = None):
    """
    Append time slice to existing zarr dataset.

    :param store: A zarr store.
    :param time_slice: Time slice to insert
    :param chunk_sizes: desired chunk sizes
    """
    if chunk_sizes:
        time_slice = chunk_dataset(time_slice, chunk_sizes, format_name='zarr')

    # Unfortunately time_slice.to_zarr(store, mode='a', append_dim='time') will replace global attributes of store
    # with attributes of time_slice (xarray bug?), which are usually empty in our case.
    # Hence, we must save our old attributes in a copy of time_slice.
    ds = zarr.open_group(store, mode='r')
    time_slice = time_slice.copy()
    time_slice.attrs.update(ds.attrs)
    if 'coordinates' in time_slice.attrs:
        # Remove 'coordinates', otherwise we get
        # ValueError: cannot serialize coordinates because the global attribute 'coordinates' already exists
        # from next time_slice.to_zarr(...) call.
        time_slice.attrs.pop('coordinates')

    time_slice.to_zarr(store, mode='a', append_dim='time')
    unchunk_dataset(store, coords_only=True) 

# 需要导入模块: import zarr [as 别名]
# 或者: from zarr import open_group [as 别名]
def load_zarr_data(zarr_fn, chrom, s1, s2, gdistkey=None):

    import zarr

    samples1 = get_sample_ids(s1)
    samples2 = get_sample_ids(s2)

    zfh = zarr.open_group(zarr_fn, mode="r")[chrom]

    samples_x = zfh["samples"][:]
    sample_name = [sid.decode() for sid in samples_x.tolist()]

    idx1 = np.array([sample_name.index(sid) for sid in samples1])
    idx2 = np.array([sample_name.index(sid) for sid in samples2])

    g = allel.GenotypeChunkedArray(zfh["calldata"]["genotype"])

    pos = allel.SortedIndex(zfh["variants"]["POS"][:])




    if gdistkey is not None:
        gdist = h5fh["variants"][gdistkey][:]
    else:
        gdist = None

    return g.take(idx1, axis=1), g.take(idx2, axis=1), pos, gdist 

# 需要导入模块: import zarr [as 别名]
# 或者: from zarr import open_group [as 别名]
def test_write_output_vars(self, in_ds, store):
        model = store.model
        model.state[("profile", "u")] = np.array([1.0, 2.0, 3.0])
        model.state[("roll", "u_diff")] = np.array([-1.0, 1.0, 0.0])
        model.state[("add", "offset")] = 2.0

        store.write_output_vars(-1, 0)

        ztest = zarr.open_group(store.zgroup.store, mode="r")

        assert ztest.profile__u.shape == (in_ds.clock.size, 3)
        np.testing.assert_array_equal(ztest.profile__u[0], np.array([1.0, 2.0, 3.0]))

        assert ztest.roll__u_diff.shape == (in_ds.out.size, 3)
        np.testing.assert_array_equal(ztest.roll__u_diff[0], np.array([-1.0, 1.0, 0.0]))

        assert ztest.add__u_diff.shape == (in_ds.out.size,)
        np.testing.assert_array_equal(
            ztest.add__u_diff, np.array([2.0, np.nan, np.nan])
        )

        # test save master clock but not out clock
        store.write_output_vars(-1, 1)
        np.testing.assert_array_equal(ztest.profile__u[1], np.array([1.0, 2.0, 3.0]))
        np.testing.assert_array_equal(
            ztest.roll__u_diff[1], np.array([np.nan, np.nan, np.nan])
        )

        # test save no-clock outputs
        store.write_output_vars(-1, -1)
        np.testing.assert_array_equal(
            ztest.profile__u_opp, np.array([-1.0, -2.0, -3.0])
        )
        assert ztest.add__offset[()] == 2.0 

# 需要导入模块: import zarr [as 别名]
# 或者: from zarr import open_group [as 别名]
def test_write_output_vars_batch(self, store_batch, model_batch1, model_batch2):
        model_batch1.state[("profile", "u")] = np.array([1.0, 2.0, 3.0])
        model_batch2.state[("profile", "u")] = np.array([4.0, 5.0, 6.0])

        model_batch1.state[("roll", "u_diff")] = np.array([-1.0, 1.0, 0.0])
        model_batch2.state[("roll", "u_diff")] = np.array([0.0, 1.0, -1.0])

        model_batch1.state[("add", "offset")] = 2.0
        model_batch2.state[("add", "offset")] = 3.0

        store_batch.write_output_vars(0, 0, model=model_batch1)
        store_batch.write_output_vars(1, 0, model=model_batch2)

        ztest = zarr.open_group(store_batch.zgroup.store, mode="r")

        assert ztest.profile__u.ndim == 3
        np.testing.assert_array_equal(
            ztest.profile__u[:, 0, :], np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
        )

        store_batch.write_output_vars(0, -1, model=model_batch1)
        store_batch.write_output_vars(1, -1, model=model_batch2)

        np.testing.assert_array_equal(ztest.add__offset[:], np.array([2.0, 3.0]))

        # test default chunk size along batch dim
        assert ztest.profile__u.chunks[0] == 1 

# 需要导入模块: import zarr [as 别名]
# 或者: from zarr import open_group [as 别名]
def test_write_index_vars(self, store):
        store.model.state[("init_profile", "x")] = np.array([1.0, 2.0, 3.0])

        store.write_index_vars()
        ztest = zarr.open_group(store.zgroup.store, mode="r")

        np.testing.assert_array_equal(ztest.x, np.array([1.0, 2.0, 3.0])) 

# 需要导入模块: import zarr [as 别名]
# 或者: from zarr import open_group [as 别名]
def test_write_index_vars_batch(self, store_batch, model_batch1):
        # ensure that no batch dim is created
        model_batch1.state[("init_profile", "x")] = np.array([1.0, 2.0, 3.0])

        store_batch.write_index_vars(model=model_batch1)
        ztest = zarr.open_group(store_batch.zgroup.store, mode="r")

        np.testing.assert_array_equal(ztest.x, np.array([1.0, 2.0, 3.0])) 

# 需要导入模块: import zarr [as 别名]
# 或者: from zarr import open_group [as 别名]
def HDF5_to_zarr(hdf5_file,CLOBBER=False,MODE=0o775):
    #-- split extension from input HDF5 file
    fileBasename,fileExtension = os.path.splitext(hdf5_file)
    #-- convert HDF5 file into zarr file
    zarr_file = '{0}.zarr'.format(fileBasename)
    #-- if zarr file exists in file system: check if HDF5 file is newer
    TEST = False
    OVERWRITE = ' (clobber)'
    #-- last modification time of HDF5 file
    hdf5_mtime = os.stat(hdf5_file).st_mtime    
    #-- check if local version of file exists
    if os.access(zarr_file, os.F_OK):
        #-- check last modification time of zarr file
        zarr_mtime = os.stat(zarr_file).st_mtime
        #-- if HDF5 file is newer: overwrite the zarr file
        if (hdf5_mtime > zarr_mtime):
            TEST = True
            OVERWRITE = ' (overwrite)'
    else:
        TEST = True
        OVERWRITE = ' (new)'

    #-- if zarr file does not exist, is to be overwritten, or CLOBBER is set
    if TEST or CLOBBER:
        #-- output string for printing files transferred
        output = '{0} -->\n\t{1}{2}\n'.format(hdf5_file,zarr_file,OVERWRITE)
        #-- copy everything from the HDF5 file to the zarr file
        with h5py.File(hdf5_file,mode='r') as source:
            dest = zarr.open_group(zarr_file,mode='w')
            #-- value checks on output zarr
            if not hasattr(dest, 'create_dataset'):
                raise ValueError('dest must be a group, got {!r}'.format(dest))
            #-- for each key in the root of the hdf5 file structure
            for k in source.keys():
                copy_from_HDF5(source[k], dest, name=k)
        #-- keep remote modification time of file and local access time
        os.utime(zarr_file, (os.stat(zarr_file).st_atime, hdf5_mtime))
        os.chmod(zarr_file, MODE)        
        #-- return the output string
        return output

#-- PURPOSE: Copy a named variable from the HDF5 file to the zarr file 

# 需要导入模块: import zarr [as 别名]
# 或者: from zarr import open_group [as 别名]
def _load(self):
        import zarr

        if self._grp is None:

            # obtain the zarr root group
            if isinstance(self._urlpath, zarr.hierarchy.Group):
                # use already-opened group, allows support for nested groups
                # as catalogs
                root = self._urlpath

            else:

                # obtain store
                if isinstance(self._urlpath, str):
                    # open store from url
                    from fsspec import get_mapper
                    store = get_mapper(self._urlpath, **self._storage_options)
                else:
                    # assume store passed directly
                    store = self._urlpath

                # open root group
                if self._consolidated:
                    # use consolidated metadata
                    root = zarr.open_consolidated(store=store, mode='r')
                else:
                    root = zarr.open_group(store=store, mode='r')

            # deal with component path
            if self._component is None:
                self._grp = root
            else:
                self._grp = root[self._component]

            # use zarr attributes as metadata
            self.metadata.update(self._grp.attrs.asdict())

        # build catalog entries
        entries = {}
        for k, v in self._grp.items():
            if isinstance(v, zarr.core.Array):
                entry = LocalCatalogEntry(name=k,
                                          description='',
                                          driver='ndzarr',
                                          args=dict(urlpath=v),
                                          catalog=self)
            else:
                entry = LocalCatalogEntry(name=k,
                                          description='',
                                          driver='zarr_cat',
                                          args=dict(urlpath=v))
            entries[k] = entry
        self._entries = entries