quadlib 0.1.0__py3-none-any.whl

quadlib/operator/sop.py

@@ -0,0 +1,250 @@
+# QuaDLib: QUAntum Dynamics LIBrary for research and methods development
+# Copyright (C) 2018-2026, The Theoretical Chemistry Group, Heidelberg University
+#
+# This software is licensed under the MIT License (https://spdx.org/licenses/MIT.html)
+
+from enum import IntEnum
+
+import numpy as np
+
+import quadlib.tools.linalg as linalg
+
+
+class OpType(IntEnum):
+    """Per-DOF operator kind in a SumOfProducts term.
+
+    UNIT: identity (no storage, no multiplication).
+    MATRIX: full N x N matrix.
+    DIAG: diagonal stored as a length-N vector.
+    """
+    UNIT = 1
+    MATRIX = 2
+    DIAG = 3
+
+
+def _is_unit_sentinel(op):
+    """A SOP-term slot is interpreted as the identity operator when the
+    user passes None, an int, or a str (legacy ``1`` / ``'1'`` sentinels).
+    The actual value is ignored.
+    """
+    return op is None or isinstance(op, (int, str))
+
+
+class SumOfProducts:
+    """Sum-of-products operator H = sum_k coeff_k prod_d O_{k,d}.
+
+    Parameters
+    ----------
+    ndims : sequence of int
+        Number of basis coefficients (primitive dimension) for each
+        degree of freedom.
+
+    Notes
+    -----
+    Operators are stored *as the user provides them*. The compiled
+    representation (``cOper``, ``cIndx``, ``cType``, ``cCoef``), built
+    on demand by :meth:`compile_operator`, is always complex; this
+    avoids silently dropping the imaginary part of operators like
+    ``PauliY`` or a momentum matrix.
+    """
+
+    def __init__(self, ndims):
+        self.ndims = np.asarray(ndims, int)
+        self.ndofs = len(ndims)
+        self.rank = 0
+        self.terms = []
+        self.coeff = []
+        self.otype = []
+        self.is_compiled = False
+
+    def add_term(self, coeff=1.0, termlist=None, termindx=None):
+        """Add one product term to the operator.
+
+        Parameters
+        ----------
+        coeff : complex or float
+            Scalar prefactor for the term.
+        termlist : sequence of operators, optional
+            Each entry is either an ndarray (1D for diagonal, 2D for
+            matrix) or a unit-operator sentinel (``None``, an int such
+            as ``1``, or the string ``'1'``).
+        termindx : sequence of int, optional
+            DOF index for each entry of ``termlist``. If omitted,
+            ``termlist`` is assumed to cover all DOFs in ascending
+            order.
+        """
+        self.is_compiled = False
+        self.coeff.append(coeff)
+        if termlist and termindx:
+            assert len(termlist) == len(termindx)
+        # Default every slot to the unit operator (stored as None).
+        term = [None] * self.ndofs
+        otype = [OpType.UNIT] * self.ndofs
+        if termlist:
+            for j, op in enumerate(termlist):
+                k = termindx[j] if termindx else j
+                if _is_unit_sentinel(op):
+                    term[k] = None
+                    otype[k] = OpType.UNIT
+                    continue
+                if len(op) != self.ndims[k]:
+                    raise ValueError('Wrong term dimension for PDOF ' + str(k))
+                term[k] = op
+                if op.ndim == 2:
+                    otype[k] = OpType.MATRIX
+                elif op.ndim == 1:
+                    otype[k] = OpType.DIAG
+                else:
+                    raise ValueError(
+                        'SOP operators must be 1D (diagonal) or 2D '
+                        '(matrix); got ndim=' + str(op.ndim)
+                    )
+        self.terms.append(term)
+        self.otype.append(otype)
+        self.rank = len(self.terms)
+
+    def multiply_term(self, k, d, vec):
+        """Apply term[k, d] to ``vec``.
+
+        ``vec`` may be a 1D array or a 2D array whose rows are the
+        SPFs; the operator is applied to the second index.
+        """
+        t = self.otype[k][d]
+        if t == OpType.UNIT:
+            return vec.copy()  # for consistency
+        if t == OpType.MATRIX:
+            return (self.terms[k][d] @ vec.T).T
+        if t == OpType.DIAG:
+            # Elementwise multiplication along the SPF axis.
+            return (np.diag(self.terms[k][d]) @ vec.T).T
+        raise ValueError('SOP otype must be UNIT, MATRIX, or DIAG')
+
+    def matrix_element_term(self, k, d, bra, ket):
+        """Compute <bra| term[k,d] |ket>.
+
+        ``bra`` and ``ket`` may be 2D arrays of SPFs stored as rows.
+        """
+        tmp = self.multiply_term(k, d, ket)
+        return bra.conj() @ tmp.T
+
+    def full_matrix_operator(self, datatype=complex):
+        """Build the explicit Hilbert-space matrix for the operator.
+
+        Defaults to ``complex`` so that operators with imaginary parts
+        (e.g. PauliY) are represented faithfully. Pass ``datatype=float``
+        only if you know every term is real.
+        """
+        N = np.prod(self.ndims)
+        H = np.zeros((N, N), datatype)
+        for k, term in enumerate(self.terms):
+            oplist = []
+            for j, op in enumerate(term):
+                t = self.otype[k][j]
+                if t == OpType.UNIT:
+                    oplist.append(np.diag(np.ones(self.ndims[j], datatype)))
+                elif t == OpType.MATRIX:
+                    oplist.append(op)
+                elif t == OpType.DIAG:
+                    oplist.append(np.diag(op))
+            h = self.coeff[k] * linalg.kron_list(oplist)
+            H = H + h
+        return H
+
+    def sparse_nnz(self):
+        """Upper bound on the non-zeros of :meth:`sparse_matrix_operator`.
+
+        Sum over terms of the product of the per-DOF non-zero counts
+        (exact per term; terms that overlap make the sum an upper bound).
+        Computed without building anything. Measured peak memory while
+        building and propagating with the sparse matrix was about 90
+        bytes per non-zero (complex128); see docs/refactor/claude_refactor_1_6.md, section 4.
+        """
+        total = 0
+        for k, term in enumerate(self.terms):
+            nnz = 1
+            for j, op in enumerate(term):
+                if self.otype[k][j] == OpType.UNIT:
+                    nnz *= int(self.ndims[j])
+                else:
+                    nnz *= int(np.count_nonzero(op))
+            total += nnz
+        return total
+
+    def sparse_matrix_operator(self, datatype=complex, format='csr'):
+        """Build the Hilbert-space matrix as a scipy sparse matrix.
+
+        Same matrix as :meth:`full_matrix_operator` (first DOF slowest,
+        C order), built term by term from Kronecker products. Unit and
+        diagonal operators stay sparse; matrix operators are stored with
+        their explicit zeros dropped, so a dense one-DOF matrix (e.g. a
+        DVR kinetic energy) contributes ``N * n_d`` non-zeros per term.
+        """
+        import scipy.sparse as sp
+        H = None
+        for k, term in enumerate(self.terms):
+            M = None
+            for j, op in enumerate(term):
+                t = self.otype[k][j]
+                if t == OpType.UNIT:
+                    m = sp.identity(self.ndims[j], dtype=datatype, format='csr')
+                elif t == OpType.MATRIX:
+                    m = sp.csr_matrix(np.asarray(op, dtype=datatype))
+                else:
+                    m = sp.diags(np.asarray(op, dtype=datatype), format='csr')
+                M = m if M is None else sp.kron(M, m, format='csr')
+            M = self.coeff[k] * M
+            H = M if H is None else H + M
+        if H is None:
+            N = int(np.prod(self.ndims))
+            H = sp.csr_matrix((N, N), dtype=datatype)
+        return H.asformat(format)
+
+    def compile_operator(self):
+        """Build the flat-array representation consumed by `quadlib.propagation.exact`.
+
+        Produces the attributes ``cCoef``, ``cOper``, ``cIndx``,
+        ``cType``. ``cOper`` is always complex-typed (see class
+        docstring).
+        """
+        nterms = len(self.coeff)
+        oplen, _ = self.len()
+        oper = np.zeros(oplen, complex)
+        indx = np.zeros((nterms, self.ndofs, 2), int)
+        i = 0
+        for k, term in enumerate(self.terms):
+            for d, op in enumerate(term):
+                nop = 0
+                t = self.otype[k][d]
+                if t == OpType.MATRIX:
+                    nop = self.ndims[d] ** 2
+                    oper[i:i + nop] = np.ravel(op)
+                elif t == OpType.DIAG:
+                    nop = self.ndims[d]
+                    oper[i:i + nop] = op
+                indx[k, d, 0] = i
+                indx[k, d, 1] = i + nop
+                i = i + nop
+        self.cCoef = np.asarray(self.coeff, dtype=complex)
+        self.cOper = oper
+        self.cIndx = indx
+        self.cType = np.asarray(self.otype, dtype=int)
+        self.is_compiled = True
+
+    def len(self):
+        """Return (total storage in flat repr, per-term storage list)."""
+        oplen = 0
+        termslen = []
+        for k, term in enumerate(self.terms):
+            tlen = 0
+            for d, _op in enumerate(term):
+                t = self.otype[k][d]
+                if t == OpType.UNIT:
+                    continue
+                if t == OpType.MATRIX:
+                    oplen += self.ndims[d] ** 2
+                    tlen += self.ndims[d] ** 2
+                elif t == OpType.DIAG:
+                    oplen += self.ndims[d]
+                    tlen += self.ndims[d]
+            termslen.append(tlen)
+        return oplen, termslen

quadlib/propagation/__init__.py

@@ -0,0 +1,17 @@
+# QuaDLib: QUAntum Dynamics LIBrary for research and methods development
+# Copyright (C) 2018-2026, The Theoretical Chemistry Group, Heidelberg University
+#
+# This software is licensed under the MIT License (https://spdx.org/licenses/MIT.html)
+
+"""
+Time propagation.
+
+integrate - generic ODE driver (scipy zvode) with output every dt
+propagate_exact - numerically exact propagation of a full tensor under
+                  a SumOfProducts operator
+propagate_cpd_A - A-vector of a WfnCPD with frozen SPFs
+"""
+
+from quadlib.propagation.cpd import propagate_cpd_A
+from quadlib.propagation.exact import propagate_exact, rhs_exact
+from quadlib.propagation.integrators import integrate

quadlib/propagation/cpd.py

@@ -0,0 +1,35 @@
+# QuaDLib: QUAntum Dynamics LIBrary for research and methods development
+# Copyright (C) 2018-2026, The Theoretical Chemistry Group, Heidelberg University
+#
+# This software is licensed under the MIT License (https://spdx.org/licenses/MIT.html)
+
+"""
+Propagation of CPD wavefunctions (quadlib.WfnCPD).
+
+Only the A-vector with frozen SPFs is supported; the coupled A/SPF
+equations are in quadlib.experimental.cpd_eom.
+"""
+
+from quadlib.propagation.integrators import integrate
+
+
+def propagate_cpd_A(wfn, dt, tfinal=None, **opts):
+    """
+    Propagate the A-vector of a WfnCPD with the SPFs frozen
+
+    S dA/dt = -i H A (WfnCPD.rhs_A), from wfn.time to tfinal with output
+    every dt (one step if tfinal is None or not beyond wfn.time). Updates
+    wfn.psi[-1] and wfn.time. opts go to quadlib.propagation.integrate.
+    Returns the integrate() output dict.
+    """
+    if tfinal is None or tfinal < wfn.time:
+        tfinal = wfn.time + dt
+
+    def funcA(t, A):
+        wfn.psi[-1] = A
+        return wfn.rhs_A()
+
+    out = integrate(funcA, wfn.psi[-1], dt, tfinal, tinit=wfn.time, **opts)
+    wfn.time = out['time'][-1]
+    wfn.psi[-1] = out['final']
+    return out

quadlib/propagation/exact.py

@@ -0,0 +1,139 @@
+# QuaDLib: QUAntum Dynamics LIBrary for research and methods development
+# Copyright (C) 2018-2026, The Theoretical Chemistry Group, Heidelberg University
+#
+# This software is licensed under the MIT License (https://spdx.org/licenses/MIT.html)
+
+"""
+Numerically exact ("standard method") propagation of a full tensor
+under a sum-of-products operator.
+"""
+
+import numpy as np
+
+import quadlib.tools.tensor as tens
+from quadlib.propagation.integrators import OutputRecorder, integrate, n_output_steps
+
+
+def propagate_exact(psi0, op, dt, tfinal, *, tinit=0.0, solver='zvode', **opts):
+    """
+    Propagate the TDSE i dpsi/dt = H psi exactly for a full tensor psi
+
+    psi0 - array of shape op.ndims (primitive dimension per DOF)
+    op - SumOfProducts operator H (time independent)
+    dt, tfinal, tinit - output interval, final and initial time
+    solver - 'zvode': integrate the TDSE with the SOP right-hand side
+                      (rhs_exact); memory ~ size of psi plus the SOP
+             'expm':  build H as a sparse matrix (op.sparse_matrix_operator)
+                      and apply exp(-i H dt) at every output step with
+                      scipy.sparse.linalg.expm_multiply; accurate to
+                      machine precision, but H must fit in memory
+    opts - store_psi, observables (both solvers; observables receive psi
+           flattened to 1D); for 'zvode' also method, atol, rtol,
+           maxstep (see quadlib.propagation.integrate)
+
+    Returns the integrate()-style output dict; 'psi' entries and 'final'
+    have the shape of psi0.
+    """
+    shape = np.asarray(psi0.shape, int)
+    if tuple(shape) != tuple(op.ndims):
+        raise ValueError('propagate_exact: psi0 shape ' + str(psi0.shape)
+                         + ' does not match operator dimensions ' + str(tuple(op.ndims)))
+    y0 = np.asarray(psi0, complex).reshape(-1)
+    if solver == 'zvode':
+        if not op.is_compiled:
+            op.compile_operator()
+        coef, oper, indx, otyp = op.cCoef, op.cOper, op.cIndx, op.cType
+
+        def dtpsi(t, A):
+            return rhs_exact(A, shape, coef, oper, indx, otyp)
+
+        out = integrate(dtpsi, y0, dt, tfinal, tinit=tinit, **opts)
+    elif solver == 'expm':
+        out = _propagate_expm(y0, op, dt, tfinal, tinit, **opts)
+    else:
+        raise ValueError("propagate_exact: solver must be 'zvode' or 'expm', got "
+                         + str(solver))
+    if 'psi' in out:
+        out['psi'] = [p.reshape(psi0.shape) for p in out['psi']]
+    out['final'] = out['final'].reshape(psi0.shape)
+    return out
+
+
+def _propagate_expm(y0, op, dt, tfinal, tinit, store_psi=True, observables=None):
+    from scipy.sparse.linalg import expm_multiply
+    H = op.sparse_matrix_operator()
+    A = -1.0j * dt * H
+    traceA = A.diagonal().sum()
+    rec = OutputRecorder(lambda t, y: np.real(np.vdot(y, H @ y)), store_psi, observables)
+    y = y0
+    rec.record(tinit, y)
+    for i in range(1, n_output_steps(tinit, tfinal, dt) + 1):
+        y = expm_multiply(A, y, traceA=traceA)
+        rec.record(tinit + i * dt, y)
+    out = rec.finish(y)
+    out['integrator'] = None
+    return out
+
+
+def rhs_exact(avector, shape, coef, oper, indx, otyp):
+    """
+    Perform the exact RHS of the TDSE
+    """
+    Nd = np.prod(shape)  # number of coefficients in tensor
+    Nf = len(shape)      # number of degrees of freedom (indices)
+    Nk = len(coef)       # number of product terms in the operator
+    ijk = np.zeros(3, int)
+    dt_a = np.zeros_like(avector)
+    for k in range(Nk): # loop operator terms
+        tmp = np.copy(avector)
+        for f in range(Nf):
+            ijk[0] = np.prod(shape[0:f])
+            ijk[1] = shape[f]
+            ijk[2] = np.prod(shape[f+1:])
+            tmp = tmp.reshape(ijk)
+            i0 = indx[k, f, 0]
+            i1 = indx[k, f, 1]
+            if otyp[k, f] == 1:    # unit operator
+                continue
+            elif otyp[k, f] == 2:  # matrix operator
+                op = oper[i0:i1].reshape(ijk[1], ijk[1])
+                tmp = np.einsum('hj,ijk->ihk', op, tmp)
+            elif otyp[k, f] == 3:  # diagonal operator
+                op = oper[i0:i1]
+                tmp = np.einsum('j,ijk->ijk', op, tmp)
+        tmp = tmp.reshape(Nd)
+        dt_a = dt_a + tmp*coef[k]
+    dt_a = -1.0j*dt_a
+    return dt_a
+
+
+def rhs_exact_alt1(avector, shape, coef, oper, indx, otyp):
+    """
+    Perform the exact RHS of the TDSE
+
+    Alternative to rhs_exact using the precomputed einsum strings in
+    quadlib.tools.tensor.esumstr; same result, not used by default.
+    """
+    Nd = np.prod(shape)  # number of coefficients in tensor
+    Nf = len(shape)      # number of degrees of freedom (indices)
+    Nk = len(coef)       # number of product terms in the operator
+    dt_a = np.zeros_like(avector)
+    for k in range(Nk): # loop operator terms
+        tmp = np.copy(avector)
+        tmp = tmp.reshape(shape)
+        for f in range(Nf):
+            i0 = indx[k, f, 0]
+            i1 = indx[k, f, 1]
+            nf = shape[f]
+            if otyp[k, f] == 1:    # unit operator
+                continue
+            elif otyp[k, f] == 2:  # matrix operator
+                op = oper[i0:i1].reshape(nf, nf)
+                tmp = np.einsum(tens.esumstr[(Nf, 2, f)], op, tmp)
+            elif otyp[k, f] == 3:  # diagonal operator
+                op = oper[i0:i1]
+                tmp = np.einsum(tens.esumstr[(Nf, 3, f)], op, tmp)
+        tmp = tmp.reshape(Nd)
+        dt_a = dt_a + tmp*coef[k]
+    dt_a = -1.0j*dt_a
+    return dt_a

quadlib/propagation/integrators.py

@@ -0,0 +1,108 @@
+# QuaDLib: QUAntum Dynamics LIBrary for research and methods development
+# Copyright (C) 2018-2026, The Theoretical Chemistry Group, Heidelberg University
+#
+# This software is licensed under the MIT License (https://spdx.org/licenses/MIT.html)
+
+import numpy as np
+from scipy.integrate import ode
+
+
+def n_output_steps(tinit, tfinal, dt):
+    """Number of output steps of size dt from tinit up to tfinal.
+
+    The tolerance absorbs float round-off in (tfinal - tinit)/dt so that
+    tfinal itself is reached when it is a multiple of dt.
+    """
+    return int(np.floor((tfinal - tinit) / dt + 1e-8))
+
+
+class OutputRecorder:
+    """Collects the output dict shared by the propagators.
+
+    energy - callable(t, psi) returning the energy expectation value
+    """
+    def __init__(self, energy, store_psi=True, observables=None):
+        self.energy = energy
+        self.store_psi = store_psi
+        self.observables = observables or {}
+        self.out = {'time': [], 'psi': [], 'norm': [], 'energy': []}
+        for name in self.observables:
+            self.out[name] = []
+
+    def record(self, t, y):
+        out = self.out
+        out['time'].append(t)
+        if self.store_psi:
+            out['psi'].append(y.copy())
+        out['norm'].append(np.real(np.vdot(y, y)))
+        out['energy'].append(self.energy(t, y))
+        for name, obs in self.observables.items():
+            out[name].append(obs(t, y))
+
+    def finish(self, y):
+        if not self.store_psi:
+            del self.out['psi']
+        self.out['final'] = y.copy()
+        return self.out
+
+
+def integrate(func, psi0, dt, tfinal=None, *,
+              tinit=0.0,
+              method='adams',
+              atol=1.e-10,
+              rtol=1.e-10,
+              maxstep=10000,
+              store_psi=True,
+              observables=None):
+    """
+    Integrate dpsi/dt = func(t, psi) with scipy's zvode, reporting every dt
+
+    One integrator instance is used for the whole interval, so its step
+    size and history carry over between output times.
+
+    func - callable dtpsi = func(t, psi), psi a 1D complex array
+    psi0 - initial state (1D)
+    dt - output interval; psi is reported at tinit, tinit + dt, ...,
+         up to and including tfinal (tfinal defaults to tinit + dt)
+    method - zvode method, 'adams' (non-stiff) or 'bdf' (stiff)
+    atol, rtol - zvode absolute and relative tolerances. The default 1e-10
+                 is deliberate: at 1e-8, the 4-mode spin-boson example
+                 (20000 states, t = 25) developed a 2.5e-3 error in the
+                 highest boson states (tests/test_spin_boson.py)
+    maxstep - maximum number of internal steps per output interval
+    store_psi - keep psi at every output time (False keeps only the last)
+    observables - optional dict {name: callable(t, psi)} evaluated at
+                  every output time
+
+    Returns a dict with keys
+        time - output times
+        psi - list of psi at each output time (only if store_psi)
+        final - psi at the last output time
+        norm - <psi|psi> at each output time
+        energy - <psi|i func(t, psi)> at each output time; the energy only
+                 if func is -iH psi with H Hermitian in an orthonormal basis
+        <name> - one list per entry of observables
+        integrator - the scipy ode instance
+    """
+    if method not in ('adams', 'bdf'):
+        raise ValueError('integrate: unknown zvode method ' + str(method))
+    if tfinal is None:
+        tfinal = tinit + dt
+    nsteps = n_output_steps(tinit, tfinal, dt)
+    r = ode(func)
+    r.set_integrator('zvode', method=method, nsteps=maxstep, atol=atol, rtol=rtol)
+    r.set_initial_value(psi0, tinit)
+    rec = OutputRecorder(lambda t, y: np.real(np.vdot(y, func(t, y)) / (-1.0j)),
+                         store_psi, observables)
+    y = np.asarray(psi0)
+    rec.record(tinit, y)
+    for i in range(1, nsteps + 1):
+        # integrate to tinit + i*dt rather than accumulating r.t + dt
+        r.integrate(tinit + i * dt)
+        if not r.successful():
+            raise RuntimeError('integrate: zvode failed at t = ' + str(r.t))
+        y = r.y
+        rec.record(r.t, y)
+    out = rec.finish(y)
+    out['integrator'] = r
+    return out

quadlib/tools/bath.py

@@ -0,0 +1,84 @@
+# QuaDLib: QUAntum Dynamics LIBrary for research and methods development
+# Copyright (C) 2018-2026, The Theoretical Chemistry Group, Heidelberg University
+#
+# This software is licensed under the MIT License (https://spdx.org/licenses/MIT.html)
+
+import numpy as np
+import scipy.integrate as spi
+
+
+def spectral_density_ohmic_bath(alpha, wcutoff):
+    """Return the spectral density function for an ohmic bath
+
+    """
+    nu = np.pi*alpha/2.0
+    def J(w):
+        return nu*w*np.exp(-w/wcutoff)
+    return J
+
+def bath_omegas_couplings(Jspectral, rho, N, wmax):
+    """Return the angular frequencies and mode couplings for a linear bath
+
+    This routine implements the equations in Section III of
+    Wang et al., J. Chem. Phys. 115, 2979 (2001)
+
+    Jspectral -- function, spectral density of the bath
+    rho -- function, mode density; int_0^{w_j} rho(w) dw = j
+    N -- number of discrete bath modes
+    wmax -- largest frequency of the bath
+
+    Returns two lists
+     [w1, w2, ...], [c1, c2, ...] as arrays
+    """
+    # Find out the discrete frequencies
+    w, r = _find_omegas(rho, N, wmax)
+    c = []
+    for j, wj in enumerate(w):
+        cj = np.sqrt(Jspectral(wj)/r[j]*2.0*wj/np.pi)
+        c.append(cj)
+    return np.asarray(w), np.asarray(c)
+
+
+def _normalize_rho(rho, N, wcut):
+    R, err = spi.quad(rho, 0, wcut)
+    fnorm = N/R
+    def rho_norm(w):
+        return rho(w)*fnorm
+    return rho_norm
+
+
+def _find_omegas(rho, N, wcut, discr=5000):
+    """
+    find the omegas for a mode density normalized to int_0^{w_j} rho(w) dw = j
+
+    rho -- function, normalized mode density
+    N -- number of modes
+    wcut -- largest omega
+    discr -- discretization for the internal numerical integration
+    """
+    # Normalize rho, just in case
+    rho_norm = _normalize_rho(rho, N, wcut)
+    # The rho_norm function is now integrated along the w-axis, yielding I(w).
+    w = [] # omegas
+    r = [] # rho[omega] (at the w points)
+    rhos = np.zeros(discr, float)
+    ws = np.linspace(0, wcut, discr)
+    dw = ws[1] - ws[0]
+    for d in range(discr):
+        rhos[d] = rho_norm(ws[d])
+    I = np.zeros(discr, float)
+    I[0] = rhos[0]*dw
+    for j in range(1, discr):
+        I[j] = rhos[j]*dw + I[j - 1]
+    # Search for the minima of (I(w) - j)**2 to locate the omegas
+    for j in range(N):
+        tmp = (I - j - 1)**2
+        k = np.argmin(tmp)
+        wj = ws[k]
+        rj = rhos[k]
+        w.append(wj)
+        r.append(rj)
+    return np.asarray(w), np.asarray(r)
+
+
+