newton_cg
=========

.. py:function:: t3toolbox.optimizers.newton_cg(geometry, kind, sample, data, x0, order = None, weight = None, verbose = False, val_sample = None, val_data = None, callback = None, use_jit = False, regularizer = None, chunk_size = 'auto', **kwargs)

   .. code-block:: python

      def newton_cg(
              geometry,                       # ragged/uniform MANIFOLD (intended) / COREWISE (must match x0)
              kind:     str,                  # 'apply' / 'entries' / 'probe' (+ '_derivatives')
              sample:   typ.Any,              # ww / index, or (ww, pp) / (index, pp) for derivatives
              data:     typ.Any,             # observed values to fit
              x0:       Point,                # initial point (zero is fine on the manifold)
              order:    typ.Optional[int] = None,
              weight:   typ.Optional[typ.Any] = None,   # residual weight ω: per-mode (probe) / ω[mode,order] (derivatives)
              verbose:  bool = False,                   # print per-iteration diagnostics (a relative-error table + more)
              val_sample: typ.Any = None,               # optional validation sample (same layout as `sample`) -> a train|val table
              val_data:   typ.Any = None,               # optional validation data (both given adds the val column)
              callback:   typ.Optional[typ.Callable] = None,  # custom callback(NewtonInfo) each iter (overrides `verbose`)
              use_jit:    bool = False,               # jit the inner CG: auto-converts x0/sample/data to jax -> a jax-backed float32 result; raises if jax absent
              regularizer: typ.Any = None,            # optional regularizer, e.g. optimizers.IdentityRegularizer(λ) (ragged only)
              chunk_size: typ.Any = 'auto',   # probe_derivatives 𝒥ᵀ memory chunk; 'auto' -> estimate_chunk_size (docs/chunking.md)
              **kwargs,                       # forwarded to backend.optimizers.newton_cg (max_newton, gtol_rel, g0norm_newton, ...)
      ) -> typ.Tuple[Point, dict]:


   Inexact Riemannian Newton-CG with an Armijo line search -- the manifold workhorse. Ragged or uniform
   ``x0`` (see :py:func:`gradient_descent`). See :py:func:`t3toolbox.backend.optimizers.newton_cg`.

   In a **warm-start continuation loop** the initial gradient norm ‖g0‖ is misleadingly small, which
   over-tightens the Newton stop and slackens CG; pass ``g0norm_newton`` / ``g0norm_cg`` (forwarded via
   ``**kwargs``) to override the reference norm the two relative stopping tests use, and
   ``cg_forcing_power`` to trade more CG iterations per Newton step for fewer Newton steps (all detailed
   on the backend).

   ``verbose=True`` prints a per-iteration diagnostic block (objective / gradient, CG stats, line search,
   and the per-``(mode, order)`` relative-error table); pass ``val_sample`` / ``val_data`` to add a
   validation column. With a ``regularizer`` attached the objective is shown split as ``obj = misfit +
   reg`` (the ``½‖ω⊙r‖²`` data misfit vs ``ρ(x)``); both parts are on every record as ``misfit`` /
   ``regularization`` (the latter ``None`` when unregularized), in ``stats['history']`` (always) and
   ``stats['diagnostics']`` (when verbose). This is a
   thin convenience over the **backend** display -- a raw-``.data`` user builds the same callback with
   :py:func:`t3toolbox.backend.optimizer_display.make_newton_display` and passes it as ``callback=`` to
   :py:func:`t3toolbox.backend.optimizers.newton_cg`. A custom ``callback`` overrides ``verbose``. Works on
   both the ragged and uniform layers (the uniform ``block_sumsq`` reduces the packed residual directly;
   validation data is packed automatically).

   For a large-``|W|`` uniform ``probe_derivatives`` fit, ``chunk_size='auto'`` (the default) sizes the
   ``𝒥ᵀ`` gradient assembly's memory automatically (via
   :py:func:`~t3toolbox.backend.sampling_derivatives.estimate_chunk_size`); pass an ``int`` / ``None`` to
   override. Other kinds and the ragged layer ignore it. See :doc:`/chunking`.

   .. rubric:: Examples

   Recover a low-rank tensor from noiseless ``apply`` measurements on the **uniform** layer -- a uniform
   ``x0`` (with ``UNIFORM_MANIFOLD``) returns a ``UniformTuckerTensorTrain`` (a zero start is fine on the
   manifold; unit-norm probe rows keep the least-squares well-conditioned):

   >>> import numpy as np
   >>> import t3toolbox.tucker_tensor_train as t3
   >>> import t3toolbox.uniform_tucker_tensor_train as ut3
   >>> import t3toolbox.uniform_manifold as ut3m
   >>> import t3toolbox.optimizers as optimizers
   >>> np.random.seed(0)
   >>> A = t3.TuckerTensorTrain.randn((6, 7, 8), (2, 2, 2), (1, 2, 2, 1))
   >>> ww = [np.random.randn(120, N) for N in (6, 7, 8)]
   >>> ww = [w / np.linalg.norm(w, axis=1, keepdims=True) for w in ww]   # unit-norm rows
   >>> b = A.apply(ww)
   >>> x0 = ut3.UniformTuckerTensorTrain.from_t3(t3.TuckerTensorTrain.zeros((6, 7, 8), (2, 2, 2), (1, 2, 2, 1)))
   >>> x_opt, stats = optimizers.newton_cg(ut3m.UNIFORM_MANIFOLD, 'apply', ww, b, x0, max_newton=30)
   >>> type(x_opt).__name__                              # returned in the same (uniform) representation
   'UniformTuckerTensorTrain'
   >>> bool(np.linalg.norm(x_opt.to_dense() - A.to_dense()) / np.linalg.norm(A.to_dense()) < 1e-6)
   True

   For a **ragged** fit pass ``t3m.MANIFOLD`` and a ``TuckerTensorTrain`` ``x0`` instead; the call is
   otherwise identical.

   **Precision — jit runs in jax's default float32.** ``use_jit=True`` opts into jax world, whose default
   dtype is float32. Fit a small exact-rank tensor three ways; the giveaway throughout is the returned
   core's dtype. First set up the problem:

   >>> import t3toolbox.manifold as t3m
   >>> np.random.seed(0)
   >>> A = t3.TuckerTensorTrain.randn((4, 5, 6), (2, 2, 2), (1, 2, 2, 1))
   >>> ww = [np.random.randn(60, N) for N in (4, 5, 6)]
   >>> ww = [w / np.linalg.norm(w, axis=1, keepdims=True) for w in ww]
   >>> b = A.probe(ww)
   >>> x0 = t3.TuckerTensorTrain.zeros((4, 5, 6), (2, 2, 2), (1, 2, 2, 1))
   >>> def rel_err(x): return float(np.linalg.norm(np.asarray(x.to_dense()) - A.to_dense()) / np.linalg.norm(A.to_dense()))

   (1) The numpy path is float64 and recovers to near machine precision (~1e-10):

   >>> x_np, _ = optimizers.newton_cg(t3m.MANIFOLD, 'probe', ww, b, x0, max_newton=20)
   >>> str(x_np.data[0][0].dtype), bool(rel_err(x_np) < 1e-8)
   ('float64', True)

   (2) ``use_jit=True`` runs the SAME code in jax's default float32 and stalls ~1000x coarser (~1e-7):

   >>> x_jit, _ = optimizers.newton_cg(t3m.MANIFOLD, 'probe', ww, b, x0, max_newton=20, use_jit=True)
   >>> str(x_jit.data[0][0].dtype), bool(rel_err(x_jit) > 1e-8)
   ('float32', True)

   (3) Enabling jax x64 restores float64 under jit -- full accuracy again (~1e-10). ``jax_enable_x64`` is
   a **global, process-wide** flag, so restore it right after the fit; the check reads values captured
   while it was on:

   >>> import jax
   >>> jax.config.update("jax_enable_x64", True)
   >>> x_x64, _ = optimizers.newton_cg(t3m.MANIFOLD, 'probe', ww, b, x0, max_newton=20, use_jit=True)
   >>> dtype_x64, ok_x64 = str(x_x64.data[0][0].dtype), bool(rel_err(x_x64) < 1e-8)
   >>> jax.config.update("jax_enable_x64", False)          # restore the default before asserting (no leak)
   >>> dtype_x64, ok_x64
   ('float64', True)