passagemath-modules 10.6.31__cp314-cp314-musllinux_1_2_x86_64.whl

sage/stats/distributions/discrete_gaussian_lattice.py

@@ -0,0 +1,908 @@
+# sage_setup: distribution = sagemath-modules
+r"""
+Discrete Gaussian Samplers over Lattices
+
+This file implements oracles which return samples from a lattice following a
+discrete Gaussian distribution. That is, if `\sigma` is big enough relative to
+the provided basis, then vectors are returned with a probability proportional
+to `\exp(-|x - c|_2^2 / (2\sigma^2))`. More precisely lattice vectors in `x \in
+\Lambda` are returned with probability:
+
+.. MATH::
+
+    \frac{\exp(-|x - c|_2^2 / (2\sigma^2))}{\sum_{v \in \Lambda} \exp(-|v|_2^2 /
+    (2\sigma^2))}.
+
+AUTHORS:
+
+- Martin Albrecht (2014-06-28): initial version
+
+- Gareth Ma (2023-09-22): implement non-spherical sampling
+
+EXAMPLES::
+
+    sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^10, 3.0)
+    sage: D(), D(), D()  # random
+    ((3, 0, -5, 0, -1, -3, 3, 3, -7, 2), (4, 0, 1, -2, -4, -4, 4, 0, 1, -4), (-3, 0, 4, 5, 0, 1, 3, 2, 0, -1))
+    sage: a = D()
+    sage: a.parent()
+    Ambient free module of rank 10 over the principal ideal domain Integer Ring
+"""
+# ******************************************************************************
+#
+#                        DGS - Discrete Gaussian Samplers
+#
+# Copyright (c) 2014, Martin Albrecht  <martinralbrecht+dgs@googlemail.com>
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+#    list of conditions and the following disclaimer.
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# The views and conclusions contained in the software and documentation are
+# those of the authors and should not be interpreted as representing official
+# policies, either expressed or implied, of the FreeBSD Project.
+# *****************************************************************************/
+
+from sage.functions.log import exp
+from sage.rings.real_mpfr import RealField
+from sage.rings.integer_ring import ZZ
+from sage.rings.rational_field import QQ
+from sage.stats.distributions.discrete_gaussian_integer import DiscreteGaussianDistributionIntegerSampler
+from sage.structure.sage_object import SageObject
+from sage.misc.cachefunc import cached_method
+from sage.misc.functional import sqrt
+from sage.misc.lazy_import import lazy_import
+from sage.misc.prandom import normalvariate
+from sage.misc.verbose import verbose
+from sage.matrix.constructor import matrix
+from sage.modules.free_module import FreeModule
+from sage.modules.free_module_element import vector
+
+lazy_import('sage.symbolic.constants', 'pi')
+
+
+def _iter_vectors(n, lower, upper, step=None):
+    r"""
+    Iterate over all integer vectors of length ``n`` between ``lower`` and
+    ``upper`` bound.
+
+    INPUT:
+
+    - ``n`` -- integer ``>0``; length
+    - ``lower`` -- integer ``< upper``; lower bound (inclusive)
+    - ``upper`` -- integer ``> lower``; upper bound (exclusive)
+    - ``step`` -- used for recursion, ignore
+
+    EXAMPLES::
+
+      sage: from sage.stats.distributions.discrete_gaussian_lattice import _iter_vectors
+      sage: list(_iter_vectors(2, -1, 2))
+      [(-1, -1), (0, -1), (1, -1), (-1, 0), (0, 0), (1, 0), (-1, 1), (0, 1), (1, 1)]
+    """
+    if step is None:
+        if ZZ(lower) >= ZZ(upper):
+            raise ValueError("expected lower < upper, but got %d >= %d" % (lower, upper))
+        if ZZ(n) <= 0:
+            raise ValueError("expected n>0 but got %d <= 0" % n)
+        step = n
+
+    assert step > 0
+    if step == 1:
+        for x in range(lower, upper):
+            v = vector(ZZ, n)
+            v[0] = x
+            yield v
+        return
+    else:
+        for x in range(lower, upper):
+            for v in _iter_vectors(n, lower, upper, step - 1):
+                v[step - 1] = x
+                yield v
+
+
+class DiscreteGaussianDistributionLatticeSampler(SageObject):
+    r"""
+    GPV sampler for Discrete Gaussians over Lattices.
+
+    EXAMPLES::
+
+        sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^10, 3.0); D
+        Discrete Gaussian sampler with Gaussian parameter σ = 3.00000000000000, c=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0) over lattice with basis
+        <BLANKLINE>
+        [1 0 0 0 0 0 0 0 0 0]
+        [0 1 0 0 0 0 0 0 0 0]
+        [0 0 1 0 0 0 0 0 0 0]
+        [0 0 0 1 0 0 0 0 0 0]
+        [0 0 0 0 1 0 0 0 0 0]
+        [0 0 0 0 0 1 0 0 0 0]
+        [0 0 0 0 0 0 1 0 0 0]
+        [0 0 0 0 0 0 0 1 0 0]
+        [0 0 0 0 0 0 0 0 1 0]
+        [0 0 0 0 0 0 0 0 0 1]
+
+
+    We plot a histogram::
+
+        sage: import warnings
+        sage: warnings.simplefilter('ignore', UserWarning)
+        sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(identity_matrix(2), 3.0)
+        sage: S = [D() for _ in range(2^12)]
+        sage: l = [vector(v.list() + [S.count(v)]) for v in set(S)]
+        sage: list_plot3d(l, point_list=True, interpolation='nn')                       # needs sage.plot
+        Graphics3d Object
+
+    REFERENCES:
+
+    - [GPV2008]_
+
+    .. automethod:: __init__
+    .. automethod:: __call__
+    """
+    @staticmethod
+    def compute_precision(precision, sigma):
+        r"""
+        Compute precision to use.
+
+        INPUT:
+
+        - ``precision`` -- integer `>= 53` nor ``None``
+        - ``sigma`` -- if ``precision`` is ``None`` then the precision of
+          ``sigma`` is used
+
+        EXAMPLES::
+
+            sage: DGL = distributions.DiscreteGaussianDistributionLatticeSampler
+            sage: DGL.compute_precision(100, RR(3))
+            100
+            sage: DGL.compute_precision(100, RealField(200)(3))
+            100
+            sage: DGL.compute_precision(100, 3)
+            100
+            sage: DGL.compute_precision(None, RR(3))
+            53
+            sage: DGL.compute_precision(None, RealField(200)(3))
+            200
+            sage: DGL.compute_precision(None, 3)
+            53
+        """
+        if precision is None:
+            try:
+                precision = ZZ(sigma.precision())
+            except AttributeError:
+                return 53
+        precision = max(53, precision)
+        return precision
+
+    def _normalisation_factor_zz(self, tau=None, prec=None):
+        r"""
+        This function returns an approximation of `\sum_{x \in B}
+        \exp(-|x|_2^2 / (2\sigma^2))`, i.e. the normalization factor such that the sum
+        over all probabilities is 1 for `B`, via Poisson summation.
+
+        INPUT:
+
+        - ``tau`` -- (default: ``None``) all vectors `v` with `|v|_2^2 \leq
+          \tau \sigma` are enumerated; if none is provided, enumerate vectors
+          with increasing norm until the sum converges to given precision. For
+          high dimension lattice, this is recommended.
+
+        - ``prec`` -- (default: ``None``) passed to :meth:`compute_precision`
+
+        EXAMPLES::
+
+            sage: # needs sage.libs.pari
+            sage: n = 3; sigma = 1.0
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^n, sigma)
+            sage: f = D.f
+            sage: nf = D._normalisation_factor_zz(); nf
+            15.7496...
+            sage: from collections import defaultdict
+            sage: counter = defaultdict(Integer)
+            sage: m = 0
+            sage: def add_samples(i):
+            ....:     global counter, m
+            ....:     for _ in range(i):
+            ....:         counter[D()] += 1
+            ....:         m += 1
+            sage: v = vector(ZZ, n, (0, 0, 0))
+            sage: v.set_immutable()
+            sage: while v not in counter:
+            ....:     add_samples(1000)
+            sage: while abs(m*f(v)*1.0/nf/counter[v] - 1.0) >= 0.1:
+            ....:     add_samples(1000)
+            sage: v = vector(ZZ, n, (-1, 2, 3))
+            sage: v.set_immutable()
+            sage: while v not in counter:
+            ....:     add_samples(1000)
+            sage: while abs(m*f(v)*1.0/nf/counter[v] - 1.0) >= 0.2:  # long time
+            ....:     add_samples(1000)
+            sage: DGL = distributions.DiscreteGaussianDistributionLatticeSampler
+            sage: D = DGL(ZZ^8, 0.5)
+            sage: D._normalisation_factor_zz(tau=3)
+            3.1653...
+            sage: D._normalisation_factor_zz()
+            6.8249...
+            sage: D = DGL(ZZ^8, 1000)
+            sage: round(D._normalisation_factor_zz(prec=100))                           # needs sage.symbolic
+            1558545456544038969634991553
+            sage: M = Matrix(ZZ, [[1, 3, 0], [-2, 5, 1], [3, -4, 2]])
+            sage: D = DGL(M, 1.7)
+            sage: D._normalisation_factor_zz()  # long time
+            Traceback (most recent call last):
+            ...
+            NotImplementedError: center must be at zero and basis must be trivial
+            sage: Sigma = Matrix(ZZ, [[5, -2, 4], [-2, 10, -5], [4, -5, 5]])
+            sage: D = DGL(ZZ^3, Sigma, [7, 2, 5])
+            sage: D._normalisation_factor_zz()                                          # needs fpylll
+            78.6804...
+            sage: M = Matrix(ZZ, [[1, 3, 0], [-2, 5, 1]])
+            sage: D = DGL(M, 3)
+            sage: D._normalisation_factor_zz()
+            Traceback (most recent call last):
+            ...
+            NotImplementedError: basis must be a square matrix
+            sage: D = DGL(ZZ^3, c=(1/2, 0, 0))
+            sage: D._normalisation_factor_zz()
+            Traceback (most recent call last):
+            ...
+            NotImplementedError: center must be at zero and basis must be trivial
+            sage: D = DGL(Matrix(3, 3, 1/2))
+            sage: D._normalisation_factor_zz()
+            Traceback (most recent call last):
+            ...
+            NotImplementedError: lattice must be integral
+        """
+        # If σ > 1:
+        # We use the Fourier transform g(t) of f(x) = exp(-k^2 / 2σ^2), but
+        # taking the norm of vector t^2 as input, and with norm_factor factored.
+        # If σ ≤ 1:
+        # The formula in docstring converges quickly since it has -1 / σ^2 in
+        # the exponent
+        def f_or_hat(x):
+            # Fun fact: If you remove this R() and delay the call to return,
+            # It might give an error due to precision error. For example,
+            # RR(1 + 100 * exp(-5.0 * pi^2)) == 0
+
+            if sigma > 1:
+                return R(exp(-pi**2 * (2 * sigma**2) * x))
+
+            return R(exp(-x / (2 * sigma**2)))
+
+        if not self.is_spherical:
+            # TODO: This is only a poor approximation placeholder.
+            # It should be easy to implement, since the Fourier transform
+            # is essentially the same, but I can't figure out how to
+            # tweak the `.qfrep` call below correctly.
+            from warnings import warn
+            warn("Note: `_normalisation_factor_zz` has not been properly "
+                 "implemented for non-spherical distributions.")
+            import itertools
+            from sage.functions.log import log
+            basis = self.B.LLL()
+            base = vector(ZZ, [v.round() for v in basis.solve_left(self._c)])
+            # BOUND is the largest integer such that |coords| <= 10^4
+            # However, this might still drift from true value for larger lattices
+            # So optimally one should fix the TODO above
+            BOUND = max(1, (self._RR(10**(4 / self.n)).ceil() - 1) // 2)
+            BOUND = min(BOUND, 10)
+            coords = itertools.product(range(-BOUND, BOUND + 1), repeat=self.n)
+            return sum(self.f((vector(u) + base) * self.B) for u in coords)
+
+        if self.B.nrows() != self.B.ncols():
+            raise NotImplementedError("basis must be a square matrix")
+
+        if self.B.base_ring() != ZZ:
+            raise NotImplementedError("lattice must be integral")
+
+        if self.is_spherical and not self._c_in_lattice_and_lattice_trivial:
+            raise NotImplementedError("center must be at zero and basis must be trivial")
+
+        sigma = self._sigma
+        prec = DiscreteGaussianDistributionLatticeSampler.compute_precision(
+            prec, sigma
+        )
+        R = RealField(prec=prec)
+        if sigma > 1:
+            det = self.B.det()
+            norm_factor = (sigma * sqrt(2 * pi))**self.n / det
+        else:
+            det = 1
+            norm_factor = 1
+
+        # qfrep computes theta series of a quadratic form, which is *half* the
+        # generating function of number of vectors with given norm (and no 0)
+        Q = self.Q
+        if tau is not None:
+            freq = Q.__pari__().qfrep(tau * sigma, 0)
+            res = R(1)
+            for x, fq in enumerate(freq):
+                res += 2 * ZZ(fq) * f_or_hat((x + 1) / det**self.n)
+            return R(norm_factor * res)
+
+        res = R(1)
+        bound = 0
+        # There might still be precision issue but whatever
+        while True:
+            bound += 1
+            cnt = ZZ(Q.__pari__().qfrep(bound, 0)[bound - 1])
+            inc = 2 * cnt * f_or_hat(bound / det**self.n)
+            if cnt > 0 and res == res + inc:
+                return R(norm_factor * res)
+            res += inc
+
+    @cached_method
+    def _maximal_r(self):
+        r"""
+        This function computes the largest value `r > 0` such that `\Sigma - r^2BB^T`
+        is positive definite.
+
+        This is equivalent to finding `\lambda_1(\Sigma / Q) = 1 / \lambda_n(Q
+        / \Sigma)`, which is done via the Power iteration method.
+
+        EXAMPLES::
+
+            sage: n = 3
+            sage: Sigma = Matrix(ZZ, [[5, -2, 4], [-2, 10, -5], [4, -5, 5]])
+            sage: c = vector(ZZ, [7, 2, 5])
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^n, Sigma, c)
+            sage: r = D._maximal_r(); r
+            0.58402...
+            sage: e_vals = (D.sigma() - r^2 * D.Q).eigenvalues()                        # needs sage.libs.pari
+            sage: assert all(e_val >= -1e-12 for e_val in e_vals)                       # needs sage.libs.pari
+        """
+        assert not self.is_spherical
+
+        Q = self.Q.change_ring(self._RR) / self._sigma.change_ring(self._RR)
+        v = Q[0].change_ring(self._RR)
+        cnt = 0
+        while cnt < 10000:
+            nv = (Q * v).normalized()
+            if (nv - v).norm() < 1e-12:
+                break
+            v = nv
+            cnt += 1
+        res = (v[0] / (Q * v)[0]).sqrt()
+        return res
+
+    def _randomise(self, v):
+        r"""
+        Randomly round to the latice coset `\ZZ + v` with Gaussian parameter
+        `r`. Used at :meth:`_call_non_spherical`.
+
+        REFERENCES:
+
+        - [Pei2010]_, Section 4.1
+
+        EXAMPLES::
+
+            sage: Sigma = Matrix(ZZ, [[5, -2, 4], [-2, 10, -5], [4, -5, 5]])
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^3, Sigma)
+            sage: all(D._randomise([0, 0, 0]).norm() <= 16 for _ in range(100))         # needs sage.symbolic
+            True
+        """
+        return vector(ZZ, [DiscreteGaussianDistributionIntegerSampler(self.r, c=vi)() for vi in v])
+
+    def __init__(self, B, sigma=1, c=0, r=None, precision=None, sigma_basis=False):
+        r"""
+        Construct a discrete Gaussian sampler over the lattice `\Lambda(B)`
+        with parameter ``sigma`` and center `c`.
+
+        INPUT:
+
+        - ``B`` -- a (row) basis for the lattice, one of the following:
+
+          - an integer matrix,
+          - an object with a ``.matrix()`` method, e.g. ``ZZ^n``, or
+          - an object where ``matrix(B)`` succeeds, e.g. a list of vectors
+
+        - ``sigma`` -- Gaussian parameter, one of the following:
+
+          - a real number `\sigma > 0` (spherical),
+          - a positive definite matrix `\Sigma` (non-spherical), or
+          - any matrix-like ``S``, equivalent to `\Sigma = SS^T`, when
+            ``sigma_basis`` is set
+
+        - ``c`` -- (default: 0) center `c`, any vector in `\ZZ^n` is
+          supported, but `c \in \Lambda(B)` is faster
+
+        - ``r`` -- (default: ``None``) rounding parameter `r` as defined in
+          [Pei2010]_; ignored for spherical Gaussian parameter; if not provided,
+          set to be the maximal possible such that `\Sigma - rBB^T` is positive
+          definite
+        - ``precision`` -- bit precision `\geq 53`
+        - ``sigma_basis`` -- boolean (default: ``False``); when set, ``sigma`` is treated as
+            a (row) basis, i.e. the covariance matrix is computed by `\Sigma = SS^T`
+
+        .. TODO::
+
+            Rename class methods like ``.f`` and hide most of them
+            (at least behind something like ``.data``).
+
+        EXAMPLES::
+
+            sage: n = 2; sigma = 3.0
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^n, sigma)
+            sage: f = D.f
+            sage: nf = D._normalisation_factor_zz(); nf                                 # needs sage.symbolic
+            56.5486677646...
+
+            sage: from collections import defaultdict
+            sage: counter = defaultdict(Integer); m = 0
+            sage: def add_samples(i):
+            ....:     global counter, m
+            ....:     for _ in range(i):
+            ....:         counter[D()] += 1
+            ....:         m += 1
+
+            sage: v = vector(ZZ, n, (-3, -3))
+            sage: v.set_immutable()
+            sage: while v not in counter: add_samples(1000)
+            sage: while abs(m*f(v)*1.0/nf/counter[v] - 1.0) >= 0.1:                     # needs sage.symbolic
+            ....:     add_samples(1000)
+
+            sage: counter = defaultdict(Integer); m = 0
+            sage: v = vector(ZZ, n, (0, 0))
+            sage: v.set_immutable()
+            sage: while v not in counter:
+            ....:     add_samples(1000)
+            sage: while abs(m*f(v)*1.0/nf/counter[v] - 1.0) >= 0.1:                     # needs sage.symbolic
+            ....:     add_samples(1000)
+
+        Spherical covariance are automatically handled::
+
+            sage: distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^3, sigma=Matrix(3, 3, 2))
+            Discrete Gaussian sampler with Gaussian parameter σ = 2.00000000000000, c=(0, 0, 0) over lattice with basis
+            <BLANKLINE>
+            [1 0 0]
+            [0 1 0]
+            [0 0 1]
+
+        The sampler supports non-spherical covariance in the form of a Gram
+        matrix::
+
+            sage: n = 3
+            sage: Sigma = Matrix(ZZ, [[5, -2, 4], [-2, 10, -5], [4, -5, 5]])
+            sage: c = vector(ZZ, [7, 2, 5])
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^n, Sigma, c)
+            sage: f = D.f
+            sage: nf = D._normalisation_factor_zz(); nf # This has not been properly implemented    # needs fpylll
+            78.6804...
+
+        We can compute the expected number of samples before sampling a vector::
+
+            sage: v = vector(ZZ, n, (11, 4, 8))
+            sage: v.set_immutable()
+            sage: 1 / (f(v) / nf)                                                                   # needs fpylll
+            2553.9461...
+
+            sage: counter = defaultdict(Integer); m = 0
+            sage: while v not in counter:
+            ....:     add_samples(1000)
+            sage: sum(counter.values())  # random
+            3000
+            sage: while abs(m*f(v)*1.0/nf/counter[v] - 1.0) >= 0.1:                     # needs fpylll sage.symbolic
+            ....:     add_samples(1000)
+
+        If the covariance provided is not positive definite, an error is thrown::
+
+            sage: Sigma = Matrix(ZZ, [[0, 1], [1, 0]])
+            sage: distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^2, Sigma)
+            Traceback (most recent call last):
+            ...
+            RuntimeError: Sigma(=[0.000000000000000  1.00000000000000]
+            [ 1.00000000000000 0.000000000000000]) is not positive definite
+
+        The sampler supports passing a basis for the covariance::
+
+            sage: n = 3
+            sage: S = Matrix(ZZ, [[2, 0, 0], [-1, 3, 0], [2, -1, 1]])
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^n, S, sigma_basis=True)
+            sage: D.sigma()
+            [ 4.00000000000000 -2.00000000000000  4.00000000000000]
+            [-2.00000000000000  10.0000000000000 -5.00000000000000]
+            [ 4.00000000000000 -5.00000000000000  6.00000000000000]
+
+        The non-spherical sampler supports offline computation to speed up
+        sampling. This will be useful when changing the center `c` is supported.
+        The difference is more significant for larger matrices. For 128x128 we
+        observe a 4x speedup (86s -> 20s)::
+
+            sage: D.offline_samples = []
+            sage: T = 2**12
+            sage: L = [D() for _ in range(T)] # 560ms
+            sage: D.add_offline_samples(T)    # 150ms
+            sage: L = [D() for _ in range(T)] # 370ms
+
+        We can also initialise with matrix-like objects::
+
+            sage: qf = matrix(3, [2, 1, 1,  1, 2, 1,  1, 1, 2])
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(qf, 3.0); D
+            Discrete Gaussian sampler with Gaussian parameter σ = 3.00000000000000, c=(0, 0, 0) over lattice with basis
+            <BLANKLINE>
+            [2 1 1]
+            [1 2 1]
+            [1 1 2]
+            sage: D().parent() is D.c().parent()                                        # needs sage.symbolic
+            True
+        """
+        precision = DiscreteGaussianDistributionLatticeSampler.compute_precision(precision, sigma)
+
+        self._RR = RealField(precision)
+        # Check if sigma is a (real) number or a scaled identity matrix
+        self.is_spherical = True
+        try:
+            self._sigma = self._RR(sigma)
+        except TypeError:
+            self._sigma = matrix(self._RR, sigma)
+            # Will it be "annoying" if a matrix Sigma has different behaviour
+            # sometimes? There should be a parameter in the constructor
+            if self._sigma == self._sigma[0, 0]:
+                self._sigma = self._RR(self._sigma[0, 0])
+            else:
+                if sigma_basis:
+                    self._sigma = self._sigma * self._sigma.T
+                if not self._sigma.is_positive_definite():
+                    raise RuntimeError(f"Sigma(={self._sigma}) is not positive definite")
+                self.is_spherical = False
+
+        # TODO: Support taking a basis for the covariance
+        try:
+            B = matrix(B)
+        except (TypeError, ValueError):
+            pass
+
+        try:
+            B = B.matrix()
+        except AttributeError:
+            pass
+
+        self.n = B.ncols()
+        self.B = B
+        self.Q = B * B.T
+        self._G = B.gram_schmidt()[0]
+        self._c_in_lattice_and_lattice_trivial = False
+
+        self.D = None
+        self.VS = None
+        self._c_mul_B_inv = None
+        self.r = r
+
+        self.set_c(c)
+
+    def _precompute_data(self):
+        r"""
+        Precompute basis data. Do not call this method directly.
+
+        EXAMPLES::
+
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^3, 3.0, c=(1,0,0))
+            sage: D.set_c((2, 0, 0))
+            sage: D
+            Discrete Gaussian sampler with Gaussian parameter σ = 3.00000000000000, c=(2, 0, 0) over lattice with basis
+            <BLANKLINE>
+            [1 0 0]
+            [0 1 0]
+            [0 0 1]
+
+        .. NOTE::
+
+           Do not call this method directly, it is called automatically from
+           :func:`DiscreteGaussianDistributionLatticeSampler.__init__`.
+        """
+
+        if self.is_spherical:
+            # deal with trivial case first, it is common
+            if self._c == 0 and self._G == 1:
+                self._c_in_lattice_and_lattice_trivial = True
+                D = DiscreteGaussianDistributionIntegerSampler(sigma=self._sigma)
+                self.D = tuple([D for _ in range(self.B.nrows())])
+                self.VS = FreeModule(ZZ, self.B.nrows())
+
+            else:
+                w = self.B.solve_left(self._c)
+                if w in ZZ ** self.B.nrows() and self._G == 1:
+                    self._c_in_lattice_and_lattice_trivial = True
+                    D = []
+                    for i in range(self.B.nrows()):
+                        sigma_ = self._sigma / self._G[i].norm()
+                        D.append(DiscreteGaussianDistributionIntegerSampler(sigma=sigma_))
+                    self.D = tuple(D)
+                    self.VS = FreeModule(ZZ, self.B.nrows())
+        else:
+            # Variables Sigma2 and r are from [Pei2010]_
+            # TODO: B is implicitly assumed to be full-rank for the
+            # non-spherical case. Remove this assumption :)
+
+            # Offline samples of B⁻¹D₁
+            self.offline_samples = []
+            self.B_inv = self.B.inverse()
+            self.sigma_inv = self._sigma.inverse()
+            self._c_mul_B_inv = self._c * self.B_inv
+
+            if self.r is None:
+                # Compute the maximal r such that (Sigma - r^2 * Q) > 0
+                self.r = self._maximal_r() * 0.9999
+            self.r = self._RR(self.r)
+
+            Sigma2 = self._sigma - self.r**2 * self.Q
+            try:
+                verbose(f"Computing Cholesky decomposition of a {Sigma2.dimensions()} matrix")
+                self.B2 = Sigma2.cholesky().T
+                self.B2_B_inv = self.B2 * self.B_inv
+            except ValueError:
+                raise ValueError("Σ₂ is not positive definite. Is your "
+                                 f"r(={self.r}) too large? It should be at most "
+                                 f"{self._maximal_r()}")
+
+    def __call__(self):
+        r"""
+        Return a new sample.
+
+        EXAMPLES::
+
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^3, 3.0, c=(1,0,0))
+            sage: L = [D() for _ in range(2^12)]
+            sage: mean_L = sum(L) / len(L)
+            sage: norm(mean_L.n() - D.c()) < 0.25
+            True
+
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^3, 3.0, c=(1/2,0,0))
+            sage: L = [D() for _ in range(2^12)]    # long time
+            sage: mean_L = sum(L) / len(L)          # long time
+            sage: norm(mean_L.n() - D.c()) < 0.25   # long time
+            True
+
+            sage: # needs numpy
+            sage: import numpy
+            sage: M = matrix(ZZ, [[1,2],[0,1]])
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(M, 20.0)
+            sage: L = [D() for _ in range(2^12)]                                               # long time
+            sage: div = numpy.mean([abs(x) for x,y in L]) / numpy.mean([abs(y) for x,y, in L]) # long time
+            sage: 0.9 < div < 1.1                                                              # long time
+            True
+
+        """
+        if not self.is_spherical:
+            v = self._call_non_spherical()
+        elif self._c_in_lattice_and_lattice_trivial:
+            v = self._call_simple()
+        else:
+            v = self._call()
+        v.set_immutable()
+        return v
+
+    def f(self, x):
+        r"""
+        Compute the Gaussian `\rho_{\Lambda, c, \Sigma}`.
+
+        EXAMPLES::
+
+            sage: Sigma = Matrix(ZZ, [[5, -2, 4], [-2, 10, -5], [4, -5, 5]])
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^3, Sigma)
+            sage: D.f([1, 0, 1])
+            0.802518797962478
+            sage: D.f([1, 0, 3])
+            0.00562800641440405
+        """
+        try:
+            x = vector(ZZ, self.n, x)
+        except TypeError:
+            try:
+                x = vector(QQ, self.n, x)
+            except TypeError:
+                x = vector(self._RR, self.n, x)
+        x -= self._c
+        if self.is_spherical:
+            return exp(-x.norm() ** 2 / (2 * self._sigma**2))
+        return exp(-x * self.sigma_inv * x / 2)
+
+    def sigma(self):
+        r"""
+        Gaussian parameter `\sigma`.
+
+        If `\sigma` is a real number, samples from this sampler will have expected norm
+        `\sqrt{n}\sigma` where `n` is the dimension of the lattice.
+
+        EXAMPLES::
+
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^3, 3.0, c=(1,0,0))
+            sage: D.sigma()
+            3.00000000000000
+        """
+        return self._sigma
+
+    def c(self):
+        r"""
+        Center `c`.
+
+        Samples from this sampler will be centered at `c`.
+
+        EXAMPLES::
+
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^3, 3.0, c=(1,0,0)); D
+            Discrete Gaussian sampler with Gaussian parameter σ = 3.00000000000000, c=(1, 0, 0) over lattice with basis
+            <BLANKLINE>
+            [1 0 0]
+            [0 1 0]
+            [0 0 1]
+
+            sage: D.c()
+            (1, 0, 0)
+        """
+        return self._c
+
+    def set_c(self, c):
+        r"""
+        Modifies center `c`.
+
+        EXAMPLES::
+
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^3, 3.0, c=(1,0,0))
+            sage: D.set_c((2, 0, 0))
+            sage: D
+            Discrete Gaussian sampler with Gaussian parameter σ = 3.00000000000000, c=(2, 0, 0) over lattice with basis
+            <BLANKLINE>
+            [1 0 0]
+            [0 1 0]
+            [0 0 1]
+        """
+        if c is None:
+            self._c = None
+            return
+
+        if c == 0:
+            c = vector(ZZ, self.n)
+        else:
+            try:
+                c = vector(ZZ, self.n, c)
+            except TypeError:
+                try:
+                    c = vector(QQ, self.n, c)
+                except TypeError:
+                    try:
+                        c = vector(self._RR, self.n, c)
+                    except TypeError:
+                        c = vector(self._RR, self.n)
+
+        self._c = c
+        self._precompute_data()
+
+    def __repr__(self):
+        r"""
+        EXAMPLES::
+
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^3, 3.0, c=(1,0,0)); D
+            Discrete Gaussian sampler with Gaussian parameter σ = 3.00000000000000, c=(1, 0, 0) over lattice with basis
+            <BLANKLINE>
+            [1 0 0]
+            [0 1 0]
+            [0 0 1]
+
+            sage: Sigma = Matrix(ZZ, [[10, -6, 1], [-6, 5, -1], [1, -1, 2]])
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^3, Sigma); D
+            Discrete Gaussian sampler with Gaussian parameter Σ =
+            [ 10.0000000000000 -6.00000000000000  1.00000000000000]
+            [-6.00000000000000  5.00000000000000 -1.00000000000000]
+            [ 1.00000000000000 -1.00000000000000  2.00000000000000], c=(0, 0, 0) over lattice with basis
+            <BLANKLINE>
+            [1 0 0]
+            [0 1 0]
+            [0 0 1]
+        """
+        if self.is_spherical:
+            sigma_str = f"σ = {self._sigma}"
+        else:
+            sigma_str = f"Σ =\n{self._sigma}"
+        return f"Discrete Gaussian sampler with Gaussian parameter {sigma_str}, c={self._c} over lattice with basis\n\n{self.B}"
+
+    def _call_simple(self):
+        r"""
+        Return a new sample assuming `c \in \Lambda(B)` and `B^* = 1`.
+
+        EXAMPLES::
+
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^3, 3.0, c=(1,0,0))
+            sage: L = [D._call_simple() for _ in range(2^12)]
+            sage: mean_L = sum(L) / len(L)
+            sage: norm(mean_L.n() - D.c()) < 0.25
+            True
+
+        .. NOTE::
+
+           Do not call this method directly, call
+           :func:`DiscreteGaussianDistributionLatticeSampler.__call__` instead.
+        """
+        w = self.VS([d() for d in self.D], check=False)
+        return w * self.B + self._c
+
+    def _call(self):
+        """
+        Return a new sample.
+
+        EXAMPLES::
+
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^3, 3.0, c=(1/2,0,0))
+            sage: L = [D._call() for _ in range(2^12)]
+            sage: mean_L = sum(L) / len(L)
+            sage: norm(mean_L.n() - D.c()) < 0.25
+            True
+
+        .. NOTE::
+
+           Do not call this method directly, call
+           :func:`DiscreteGaussianDistributionLatticeSampler.__call__` instead.
+        """
+        v = 0
+        c, sigma, B = self._c, self._sigma, self.B
+
+        m = self.B.nrows()
+
+        for i in range(m - 1, -1, -1):
+            b_ = self._G[i]
+            c_ = c.dot_product(b_) / b_.dot_product(b_)
+            sigma_ = sigma / b_.norm()
+            assert sigma_ > 0
+            z = DiscreteGaussianDistributionIntegerSampler(sigma=sigma_, c=c_, algorithm='uniform+online')()
+            c = c - z * B[i]
+            v = v + z * B[i]
+        return v
+
+    def add_offline_samples(self, cnt=1):
+        """
+        Precompute samples from `B^{-1}D_1` to be used in :meth:`_call_non_spherical`.
+
+        EXAMPLES::
+
+            sage: Sigma = Matrix([[5, -2, 4], [-2, 10, -5], [4, -5, 5]])
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^3, Sigma)
+            sage: assert not D.is_spherical
+            sage: D.add_offline_samples(2^12)
+            sage: L = [D() for _ in range(2^12)] # Takes less time
+        """
+        # Just to document the difference with [Pei2010]_, in the paper (Algo 1)
+        # he samples from Λ + c, but we instead sample from Λ with distribution
+        # sampled at c (D_{Λ, c}), but that's the same as c + D_{Λ - c}
+        # Also, we use row notation instead of column notation. Sorry.
+        for _ in range(cnt):
+            coord = [normalvariate(mu=0, sigma=1) for _ in range(self.n)]
+            self.offline_samples.append(vector(self._RR, coord) * self.B2_B_inv)
+
+    def _call_non_spherical(self):
+        """
+        Return a new sample.
+
+        EXAMPLES::
+
+            sage: Sigma = Matrix([[5, -2, 4], [-2, 10, -5], [4, -5, 5]])
+            sage: D = distributions.DiscreteGaussianDistributionLatticeSampler(ZZ^3, Sigma, c=(1/2,0,0))
+            sage: L = [D._call_non_spherical() for _ in range(2^12)]
+            sage: mean_L = sum(L) / len(L)
+            sage: norm(mean_L.n() - D.c()) < 0.25
+            True
+
+        .. NOTE::
+
+           Do not call this method directly, call
+           :func:`DiscreteGaussianDistributionLatticeSampler.__call__` instead.
+        """
+        if len(self.offline_samples) == 0:
+            self.add_offline_samples()
+        vec = self._c_mul_B_inv - self.offline_samples.pop()
+        return self._randomise(vec) * self.B