Skip to content

Commit

Permalink
Improve dask array name determinism
Browse files Browse the repository at this point in the history
  • Loading branch information
sjperkins committed Mar 18, 2024
1 parent 763a745 commit 52b0939
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 29 deletions.
31 changes: 2 additions & 29 deletions xarrayfits/fits.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

"""Main module."""

from functools import partial, reduce
from functools import reduce
from itertools import product
import logging
import os
Expand Down Expand Up @@ -38,32 +38,6 @@ def short_fits_file(table_name):
return os.path.split(table_name.rstrip(os.sep))[1]


def fits_open_graph(fits_file, **kwargs):
"""
Generate a dask graph containing fits open commands
Parameters
----------
fits_file : str
FITS filename
**kwargs (optional) :
Keywords arguments passed to the :meth:`astropy.io.fits.open`
command.`
Returns
-------
tuple
Graph key associated with the opened file
dict
Dask graph containing the graph open command
"""
token = dask.base.tokenize(fits_file, kwargs)
fits_key = ("open", short_fits_file(fits_file), token)
fits_graph = {fits_key: (partial(FitsProxy, **kwargs), fits_file)}
return fits_key, fits_graph


BITPIX_MAP = {
np.int8: 8,
np.int16: 16,
Expand Down Expand Up @@ -109,9 +83,8 @@ def generate_slice_gets(fits_proxy, hdu, shape, dtype, chunks):
with the ``hdu``.
"""

token = dask.base.tokenize(fits_proxy)
token = dask.base.tokenize(fits_proxy, hdu, dtype)
name = "-".join((short_fits_file(fits_proxy._filename), "slice", token))

dsk_chunks = da.core.normalize_chunks(chunks, shape)

# Produce keys and slices
Expand Down
3 changes: 3 additions & 0 deletions xarrayfits/fits_proxy.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,5 +57,8 @@ def hdu(self):
self._hdul = fits.open(self._filename, **self._kwargs)
return self._hdul

def __hash__(self):
return hash((self._filename, tuple(set(self._kwargs.items()))))

def __reduce__(self):
return (FitsProxy.from_reduce_args, (self._filename, self._kwargs))

0 comments on commit 52b0939

Please sign in to comment.