Generates docs for ndarrays

hail/python/hail/conftest.py

@@ -3,6 +3,7 @@
 import pytest
 import shutil
 import tempfile
+import numpy as np
 
 import hail as hl
 
@@ -43,6 +44,7 @@ def init(doctest_namespace):
 
 def generate_datasets(doctest_namespace):
     doctest_namespace['hl'] = hl
+    doctest_namespace['np'] = np
 
     ds = hl.import_vcf('data/sample.vcf.bgz')
     ds = ds.sample_rows(0.03)

hail/python/hail/docs/api.rst

@@ -34,6 +34,7 @@ Modules
     aggregators <aggregators>
     scans <scans>
     methods <methods/index>
+    nd <nd/index>
     utils <utils/index>
     linalg <linalg/index>
     stats <stats/index>

hail/python/hail/docs/expressions.rst

@@ -29,4 +29,5 @@ Expressions
     StringExpression
     StructExpression
     TupleExpression
-
+    NDArrayExpression
+    NDArrayNumericExpression

hail/python/hail/docs/nd/index.rst

@@ -0,0 +1,29 @@
+nd
+========
+
+
+.. toctree::
+    :maxdepth: 2
+
+.. currentmodule:: hail.nd
+
+.. rubric:: NDArray Functions
+
+Notes
+_____
+This is a recently added, experimental module. We would love to hear what use cases you have for this as we expand this functionality
+
+.. autosummary::
+
+    array
+    arange
+    full
+    zeros
+    ones
+
+.. autofunction:: array
+.. autofunction:: arange
+.. autofunction:: full
+.. autofunction:: zeros
+.. autofunction:: ones
+

hail/python/hail/docs/overview/expressions.rst

@@ -107,7 +107,7 @@ or :mod:`scipy`.
 Hail has many subclasses of :class:`.Expression` -- one for each Hail type. Each
 subclass has its own constructor method. For example, if we have a list of Python
 integers, we can convert this to a Hail :class:`.ArrayNumericExpression` with
-:func:`.array`:
+:func:`~hail.expr.functions.array`:
 
     >>> a = hl.array([1, 2, -3, 0, 5])
     >>> a
@@ -128,7 +128,7 @@ Hail arrays can be indexed and sliced like Python lists or :mod:`numpy` arrays:
     >>> a[1:-1]
     <ArrayNumericExpression of type array<int32>>
 
-In addition to constructor methods like :func:`.array` and :func:`.bool`,
+In addition to constructor methods like :func:`~hail.expr.functions.array` and :func:`.bool`,
 Hail expressions can also be constructed with the :func:`.literal` method,
 which will impute the type of of the expression.
 

hail/python/hail/expr/expressions/typed_expressions.py

@@ -3453,8 +3453,9 @@ def T(self):
 
     @typecheck_method(axes=nullable(tupleof(int)))
     def transpose(self, axes=None):
-        """Permute the dimensions of this ndarray according to the ordering of `axes`. Axis `j` in the `i`th index of
-        `axes` maps the `j`th dimension of the ndarray to the `i`th dimension of the output ndarray.
+        """
+        Permute the dimensions of this ndarray according to the ordering of axes. Axis j in the ith index of
+        axes maps the jth dimension of the ndarray to the ith dimension of the output ndarray.
 
         Parameters
         ----------
@@ -3504,9 +3505,9 @@ def shape(self):
         shape_type = ttuple(*[tint64 for _ in range(self.ndim)])
         return construct_expr(NDArrayShape(self._ir), shape_type, self._indices, self._aggregations)
 
-    opt_long_slice_ = sliceof(nullable(expr_int64), nullable(expr_int64), nullable(expr_int64))
+    _opt_long_slice = sliceof(nullable(expr_int64), nullable(expr_int64), nullable(expr_int64))
 
-    @typecheck_method(item=oneof(expr_int64, opt_long_slice_, tupleof(oneof(expr_int64, opt_long_slice_))))
+    @typecheck_method(item=oneof(expr_int64, _opt_long_slice, tupleof(oneof(expr_int64, _opt_long_slice))))
     def __getitem__(self, item):
         if not isinstance(item, tuple):
             item = (item,)
@@ -3782,76 +3783,6 @@ def __matmul__(self, other):
 
         return res if result_ndim > 0 else res[()]
 
-    @typecheck_method(axis=nullable(oneof(int, sequenceof(int))))
-    def sum(self, axis=None):
-        """Sum along one or more dimensions of the ndarray. If no axes are given, the entire NDArray will
-        be summed to a single `NumericExpression`.
-
-        Parameters
-        ----------
-        axis : :obj: `int` or :obj: `list` of :obj: `int:, optional
-
-        Returns
-        -------
-        :class:`.NDArrayNumericExpression`
-        """
-        if axis is None:
-            axes = list(range(self.ndim))
-        else:
-            axes = wrap_to_list(axis)
-
-        for axis in axes:
-            if not 0 <= axis <= self.ndim:
-                raise ValueError(f'Invalid axis {axis}. Axis must be between 0 and {self.ndim}.')
-
-        if len(set(axes)) != len(axes):
-            raise ValueError(f'Axes should not be repeated: {axes}')
-
-        return construct_expr(NDArrayAgg(self._ir, axes),
-                              tndarray(self._type.element_type, self.ndim - len(axes)),
-                              self._indices,
-                              self._aggregations)
-
-    @typecheck_method(uri=str)
-    def save(self, uri):
-        """Write out the NDArray to the given path as in .npy format. If the URI does not
-        end with ".npy" the file extension will be appended. This method reflects the numpy
-        `save` method. NDArrays saved with this method can be loaded into numpy using numpy
-        `load`.
-
-        Examples
-        --------
-        >>> import numpy as np
-        >>> nd.save('file://local/file') # doctest: +SKIP
-        >>> np.load('/local/file.npy') # doctest: +SKIP
-        array([[1, 2],
-               [3, 4]], dtype=int32)
-
-        Parameters
-        ----------
-        uri : :obj: `str`
-        """
-        if not uri.endswith('.npy'):
-            uri += '.npy'
-
-        Env.backend().execute(NDArrayWrite(self._ir, hl.str(uri)._ir))
-
-    def to_numpy(self):
-        """Execute and convert this NDArray to a `NumPy` ndarray.
-
-        Examples
-        --------
-        >>> a = nd.to_numpy() # doctest: +SKIP
-
-        Returns
-        -------
-        :class:`numpy.ndarray`
-        """
-        # FIXME Use filesystem abstraction instead when that is ready
-        temp_file = tempfile.NamedTemporaryFile(suffix='.npy').name
-        self.save(temp_file)
-        return np.load(temp_file)
-
 
 scalars = {tbool: BooleanExpression,
            tint32: Int32Expression,

hail/python/hail/expr/types.py

@@ -696,7 +696,7 @@ class tarray(HailType):
     See Also
     --------
     :class:`.ArrayExpression`, :class:`.CollectionExpression`,
-    :func:`.array`, :ref:`sec-collection-functions`
+    :func:`~hail.expr.functions.array`, :ref:`sec-collection-functions`
     """
 
     @typecheck_method(element_type=hail_type)

hail/python/hail/nd/nd.py

@@ -9,6 +9,35 @@
 
 
 def array(input_array):
+    """Construct an :class:`.NDArrayExpression`
+
+    Examples
+    --------
+
+    >>> hl.eval(hl._nd.array([1, 2, 3, 4]))
+    array([1, 2, 3, 4], dtype=int32)
+
+    >>> hl.eval(hl._nd.array([[1, 2, 3], [4, 5, 6]]))
+    array([[1, 2, 3],
+       [4, 5, 6]], dtype=int32)
+
+    >>> hl.eval(hl._nd.array(np.identity(3)))
+    array([[1., 0., 0.],
+       [0., 1., 0.],
+       [0., 0., 1.]])
+
+    >>> hl.eval(hl._nd.array(hl.range(10, 20)))
+    array([10, 11, 12, 13, 14, 15, 16, 17, 18, 19], dtype=int32)
+
+    Parameters
+    ----------
+    input_array : :class:`.ArrayExpression` or numpy ndarray or nested python lists
+
+    Returns
+    -------
+    :class:`.NDArrayExpression`
+        An ndarray based on the input array.
+    """
     return _ndarray(input_array)
 
 
@@ -47,12 +76,40 @@ def arange(start, stop=None, step=1) -> NDArrayNumericExpression:
     Returns
     -------
     :class:`.NDArrayNumericExpression`
+        A 1-dimensional ndarray from `start` to `stop` by `step`.
     """
     return array(hl.range(start, stop, step))
 
 
 @typecheck(shape=oneof(expr_int64, tupleof(expr_int64), expr_tuple()), value=expr_any, dtype=nullable(HailType))
 def full(shape, value, dtype=None):
+    """Creates a hail :class:`.NDArrayNumericExpression` full of the specified value.
+
+    Examples
+    --------
+
+    Create a 5 by 7 NDArray of type :py:data:`.tfloat64` 9s.
+
+    >>> hl._nd.full((5, 7), 9)
+
+    It is possible to specify a type other than :py:data:`.tfloat64` with the `dtype` argument.
+
+    >>> hl._nd.full((5, 7), 9, dtype=hl.tint32)
+
+    Parameters
+    ----------
+    shape : `tuple` or :class:`.TupleExpression`
+            Desired shape.
+    value : :class:`.Expression` or python value
+            Value to fill ndarray with.
+    dtype : :class:`.HailType`
+            Desired hail type.
+
+    Returns
+    -------
+    :class:`.NDArrayNumericExpression`
+        An ndarray of the specified shape filled with the specified value.
+    """
     if isinstance(shape, Int64Expression):
         shape_product = shape
     else:
@@ -67,26 +124,30 @@ def zeros(shape, dtype=hl.tfloat64):
        Examples
        --------
 
-       Create a 5 by 7 NDArray of type `tfloat64` zeros.
+       Create a 5 by 7 NDArray of type :py:data:`.tfloat64` zeros.
 
        >>> hl._nd.zeros((5, 7))
 
-       It is possible to specify a type other than `tfloat64` with the `dtype` argument.
+       It is possible to specify a type other than :py:data:`.tfloat64` with the `dtype` argument.
 
        >>> hl._nd.zeros((5, 7), dtype=hl.tfloat32)
 
 
        Parameters
        ----------
-       shape : `tuple` or `TupleExpression`
+       shape : `tuple` or :class:`.TupleExpression`
             Desired shape.
-       dtype : `HailType`
+       dtype : :class:`.HailType`
             Desired hail type.
 
+       See Also
+       --------
+       :func:`.full`
+
        Returns
        -------
        :class:`.NDArrayNumericExpression`
-           NDArray of the specified size full of zeros.
+           ndarray of the specified size full of zeros.
        """
     return full(shape, 0, dtype)
 
@@ -98,25 +159,30 @@ def ones(shape, dtype=hl.tfloat64):
        Examples
        --------
 
-       Create a 5 by 7 NDArray of type `tfloat64` ones.
+       Create a 5 by 7 NDArray of type :py:data:`.tfloat64` ones.
 
        >>> hl._nd.ones((5, 7))
 
-       It is possible to specify a type other than `tfloat64` with the `dtype` argument.
+       It is possible to specify a type other than :py:data:`.tfloat64` with the `dtype` argument.
 
        >>> hl._nd.ones((5, 7), dtype=hl.tfloat32)
 
 
        Parameters
        ----------
-       shape : `tuple` or `TupleExpression`
+       shape : `tuple` or :class:`.TupleExpression`
             Desired shape.
-       dtype : `HailType`
+       dtype : :class:`.HailType`
             Desired hail type.
 
+
+       See Also
+       --------
+       :func:`.full`
+
        Returns
        -------
        :class:`.NDArrayNumericExpression`
-           NDArray of the specified size full of ones.
+           ndarray of the specified size full of ones.
        """
     return full(shape, 1, dtype)