#!/usr/bin/env python
# coding: utf-8

# (param-model-nb)=
# # Parameterized `TBModel`
# 
# This notebook shows how to build **parameterized** tight-binding models in PythTB, how to evaluate them for different parameter values, and how to freeze those parameters into the model permanently.
# 
# :::{admonition} Remember
# :class: note
# 
# Throughout, keep in mind that very parameter value is consumed as a scalar (or a 1‑D array of scalars for sweeps).
# 
# :::
# 
# :::{admonition} What you will learn
# :class: tip
# 
# 1. Define symbolic hopping/onsite terms.
# 2. Explore parameter sweeps in evaluation routines (Hamiltonian, eigenvalues, etc.).
# 3. Permanently “freeze” parameters with `TBModel.set_parameters` and the copy-returning
#    `TBModel.with_parameters`.
# 4. Extend to spinful models, demonstrating how to supply $2\times 2$ matrix blocks parametrically.
# 
# ::: 

# In[ ]:


import numpy as np
from pythtb import TBModel, Lattice


# ## Spinless model with a symbolic hopping
# 
# We begin with a familiar two-orbital spinless model on a honeycomb-like lattice. The nearest-neighbor hopping is parameterised by a single symbol `t`.

# In[ ]:


lat_vecs = [[1.0, 0.0], [0.5, np.sqrt(3.0) / 2.0]]
orb_vecs = [[0.0, 0.0], [1.0 / 3.0, 1.0 / 3.0]]

lattice = Lattice(lat_vecs, orb_vecs, periodic_dirs=...)
tb = TBModel(lattice)

# Fixed on-site energies
tb.set_onsite([0.0, 0.0])

# Symbolic nearest-neighbour hopping: string 't'
tb.set_hop("t", 0, 1, [0, 0])
tb.set_hop(0.3, 1, 0, [1, 0])
tb.set_hop(0.1, 1, 0, [0, 1])


# At this stage, the model knows it depends on a parameter named `t`. Notice in the printout below that the hopping from site 0 to site 1 with lattice vector `[0, 0]` is given as `t`, while the other hoppings are numeric values.

# In[ ]:


tb.parameters


# In[ ]:


print(tb)


# If we try to, for example, build the Hamiltonian without providing `t`, PythTB raises an error because the symbol has no value.

# In[ ]:


try:
    tb.hamiltonian()
except ValueError as exc:
    print(exc)


# ### Providing parameter values during evaluation
# 
# Every evaluation routine (`hamiltonian`, `solve_ham`, `velocity`, etc.) accepts keyword arguments for parameters. Values can be **scalars** or **one-dimensional arrays of scalars**.

# In[ ]:


k_mesh = tb.k_uniform_mesh([20, 20])

# Single value (scalar)
H_single = tb.hamiltonian(k_mesh, t=0.5)

# Sweep over five scalar values (1-D array of scalars)
t_values = np.linspace(-1.0, 1.0, 5)
H_sweep = tb.hamiltonian(k_mesh, t=t_values)

print("Hamiltonian shape with t=0.5:")
print(H_single.shape)
print("Hamiltonian shape sweeping over 5 t values:")
print(H_sweep.shape)


# The sweep result has the parameter axis tacked on after the k-point axis. Here, we have 5 values of `t` for each of the 400 k-points, resulting in a Hamiltonian array of shape `(400, 5, 2, 2)`.
# 
# `solve_ham` works with the same pattern.

# In[ ]:


evals_single = tb.solve_ham(k_mesh, t=0.5)  # passing a single scalar value
evals_sweep = tb.solve_ham(k_mesh, t=t_values)  # passing a 1-D array of scalars

print("Energies shape with t=0.5:")
print(evals_single.shape)
print("Energies shape sweeping over 5 t values:")
print(evals_sweep.shape)


# The `velocity` method also accepts parameter values in the same way. An important caveat is when providing parameter sweeps: the derivative of the Hamiltonian with respect to the range of parameters is computed via finite differences and appended to the leading axis alongside the k-derivatives. 
# 
# :::{admonition} Notice below
# 
# The velocity array has shape `(2, 400, 2, 2)` when passing a scalar `t`, but shape `(3, 400, 5, 2, 2)` when passing a sweep of 5 `t` values. An extra axis for the parameter values has been added after the k-values, and the finite-difference derivative with respect to `t` has been tacked onto the front axis.
# 
# :::

# In[ ]:


velocity_single = tb.velocity(k_mesh, t=0.5)
velocity_sweep = tb.velocity(k_mesh, t=t_values)

print("Velocity shape with t=0.5:")
print(velocity_single.shape)
print("Velocity shape sweeping over 5 t values:")
print(velocity_sweep.shape)


# ### Freezing parameters in place
# 
# `TBModel.set_parameters` replaces symbolic providers (strings or callables) with the numeric value you supply. After calling it, those terms become permanent numbers.
# 
# :::{important}
# 
# Every parameter passed to `set_parameters` must be a scalar. Arrays of scalars are only supported in the evaluation routines, not for freezing.
# 
# :::

# In[ ]:


tb.set_parameters(t=0.8)
tb.parameters


# In[ ]:


print(tb)


# In[ ]:


# Once frozen, the model no longer expects 't'
H_frozen = tb.hamiltonian(k_mesh)
print("Hamiltonian shape after freezing parameters:")
print(H_frozen.shape)


# If you prefer not to modify the original model, use `TBModel.with_parameters`. It returns a copy with the specified parameters frozen.

# In[ ]:


tb_symbolic = TBModel(lattice)
tb_symbolic.set_onsite([0.0, 0.0])
tb_symbolic.set_hop("t", 0, 1, [0, 0])
tb_symbolic.set_hop(0.3, 1, 0, [1, 0])
tb_symbolic.set_hop(0.1, 1, 0, [0, 1])

tb_frozen = tb_symbolic.with_parameters(t=0.3)

# tb_symbolic still demands a parameter, tb_frozen does not
try:
    tb_symbolic.hamiltonian(k_mesh)
except ValueError as exc:
    print("symbolic:", exc)

H_frozen = tb_frozen.hamiltonian(k_mesh)
print("Frozen Hamiltonian shape:")
print(H_frozen.shape)


# ## Spinless model with multiple symbolic terms
# 
# This time, we create a model with multiple symbolic hopping terms. The parameters may take the same name or different names.
# Those with the same name will share the same value when evaluated. Below, both orbitals share the onsite parameter `m`, and two hoppings have different parameter names `t1` and `t2`.

# In[ ]:


lat_vecs = [[1, 0], [0, 1]]
orb_vecs = [[0, 0], [1 / 2, 1 / 2]]
lat = Lattice(lat_vecs=lat_vecs, orb_vecs=orb_vecs, periodic_dirs=...)
model = TBModel(lattice=lat, spinful=False)

model.set_onsite("m", ind_i=0)
model.set_onsite("m", ind_i=1)
model.set_hop("t1", 0, 1, [0, 0])
model.set_hop("t2", 1, 0, [0, 1])
print(model)


# Below we will set two of the parameters to specific values, and leave `t2` as a free parameter to be specified during evaluation. 

# In[ ]:


model.set_parameters(m=1, t1=0.8)

print(model)


# If we want to reinstate a parameter as symbolic after freezing it, we can use the same `set_onsite` and `set_hop` methods as when building the model initially. Here, we set the onsite energies back to depend on the parameter `m`.

# In[ ]:


model.set_onsite("m", ind_i=0)
model.set_onsite("m", ind_i=1)


# Now we can pass `m` as a varying parameter during evaluation and set `t2` to a fixed value.

# In[ ]:


H = model.hamiltonian(k_mesh, m=np.linspace(0, 1, 20), t2=0.4)

print("Hamiltonian shape with varying m:")
print(H.shape)


# Or, we can pass both `m` and `t2` as varying parameters.

# In[ ]:


H = model.hamiltonian(k_mesh, m=np.linspace(0, 1, 20), t2=np.linspace(0.1, 0.5, 10))

print("Hamiltonian shape with varying m and t2:")
print(H.shape)


# The shape of the Hamiltonian (or any other observables that accept parameter sweeps) will reflect the order in which the varying parameters are provided. Here, `m` varies over 20 values and `t2` over 10 values, resulting in a Hamiltonian array of shape `(400, 20, 10, 2, 2)`. Instead, if we had provided `t2` first, the shape would be `(400, 10, 20, 2, 2)`.

# In[ ]:


H = model.hamiltonian(k_mesh, t2=np.linspace(0, 1, 10), m=np.linspace(0.1, 0.5, 20))

print("Hamiltonian shape with varying t2 and m:")
print(H.shape)


# ## Spinless model with callable parameters

# If we want more control over how parameters are evaluated, we can use callables instead of strings when building the model. Below, we create a model similar to the first example, but with a hopping term defined by a function of `t`.
# 
# :::{tip}
# 
# A convenient way to define simple parameter functions is to use `lambda` functions, as shown below. `lambda` functions are anonymous functions defined in a single line. The syntax is `lambda arguments: expression`. The names of the arguments correspond to the parameter names used when evaluating the model.
# 
# :::

# In[ ]:


lat_vecs = [[1.0, 0.0], [0.5, np.sqrt(3.0) / 2.0]]
orb_vecs = [[0.0, 0.0], [1.0 / 3.0, 1.0 / 3.0]]

lattice = Lattice(lat_vecs, orb_vecs, periodic_dirs=...)
tb = TBModel(lattice)

# Fixed on-site energies
tb.set_onsite([0.0, 0.0])

# Hopping with a callable parameter
tb.set_hop(lambda t: np.cos(t), 0, 1, [0, 0])
tb.set_hop(0.3, 1, 0, [1, 0])
tb.set_hop(0.1, 1, 0, [0, 1])
print(tb)


# Everything else works the same way as with string parameters. We specify value(s) for the parameter `t` when evaluating the model.

# In[ ]:


tb.set_parameters(t=np.pi)
print(tb)


# We can do more complicated things with callables, such as using mutliple parameters for a single callable, some of which may be shared with other symbolic terms. Here, we create a model with two symbolic hoppings defined by callables that share parameters `r`, `theta` and `phi`.

# In[ ]:


lat_vecs = [[1.0, 0.0], [0.5, np.sqrt(3.0) / 2.0]]
orb_vecs = [[0.0, 0.0], [1.0 / 3.0, 1.0 / 3.0]]

lattice = Lattice(lat_vecs, orb_vecs, periodic_dirs=...)
tb = TBModel(lattice)

# Fixed on-site energies
tb.set_onsite([0.0, 0.0])

# Hopping with a callable parameter
tb.set_hop(lambda r, theta, phi: r * np.sin(theta) * np.sin(phi), 0, 1, [0, 0])
tb.set_hop(lambda r, theta, phi: r * np.sin(theta) * np.cos(phi), 1, 0, [1, 0])
tb.set_hop(lambda r, theta: r * np.cos(theta), 1, 0, [0, 1])
print(tb)


# In[ ]:


tb.set_parameters(r=0.5, theta=np.pi / 4, phi=np.pi / 3)
print(tb)


# ## Spinful model with parameterized $2\times 2$ blocks
# 
# Spinful on-site and hopping blocks are $2\times 2$ Hermitian matrices. You can still parameterize them, but each matrix needs to be built from scalar coefficients.
# 
# Trying to pass a matrix directly into `set_parameters` fails.

# In[ ]:


lat_spin = Lattice([[1.0]], [[0.0]], periodic_dirs=[0])
spin_tb = TBModel(lat_spin, spinful=True)

spin_tb.set_hop("t_spin", 0, 0, [1])

try:
    spin_tb.set_parameters(t_spin=np.array([[0.1, 0.2j], [-0.2j, -0.1]]))
except TypeError as exc:
    print(exc)


# Instead, we would need to use a callable that sets the desired matrix element symbolically. For example, we can define a spinful onsite block with parameters `a` and `b` as follows:

# In[ ]:


lat_spin = Lattice([[1.0]], [[0.0]], periodic_dirs=[0])
spin_tb = TBModel(lat_spin, spinful=True)

spin_tb.set_hop(lambda a, b: np.array([[a, b], [np.conjugate(b), a]]), 0, 0, [1])


# We could also only parameterize some of the matrix elements, while keeping others fixed. Here, we define a hopping block with fixed diagonal elements and parameterized off-diagonal elements.

# In[ ]:


lat_spin = Lattice([[1.0]], [[0.0]], periodic_dirs=[0])
spin_tb = TBModel(lat_spin, spinful=True)

spin_tb.set_hop(lambda b: np.array([[0.1, b], [np.conjugate(b), -0.1]]), 0, 0, [1])


# The key point is that parameters represent scalar values, so the callable must return a something consistent with what `set_hop` and `set_onsite` expect: either a scalar (for spinless) or a $2\times 2$ Hermitian matrix or Pauli vector (for spinful).

# ### Pauli decomposition pattern
# 
# Expose scalar coefficients for the Pauli basis (identity plus $\sigma_x$, $\sigma_y$, $\sigma_z$).

# In[ ]:


spin_tb = TBModel(lat_spin, spinful=True)

spin_tb.set_hop(
    lambda t0, t1, t2, t3: [
        t0,
        t1,
        t2,
        t3,
    ],  # internally converts Pauli vector to 2x2 matrix
    0,
    0,
    [1],
)

spin_tb.set_parameters(t0=0, t1=1, t2=0, t3=1)
print(spin_tb)


# ### Solver calls and sweeps
# 
# Evaluation routines also insist on scalar inputs (arrays for sweeps). Below we sweep `t0` while keeping the other coefficients fixed; the callable assembles the matrix internally for each scalar.

# In[ ]:


spin_tb = TBModel(lat_spin, spinful=True)

spin_tb.set_hop(
    lambda t0, t1, t2, t3: [t0, t1, t2, t3],
    0,
    0,
    [1],
)

t0_vals = np.linspace(-0.2, 0.2, 5)
H_spin = spin_tb.hamiltonian(
    k_pts=[[0], [0.5], [1]], t0=t0_vals, t1=0.05, t2=0.0, t3=-0.10
)
H_spin.shape  # (Nk, len(t0_vals), norb, nspin, norb, nspin)


# ## Summary
# 
# - Every parameter value—whether supplied to `hamiltonian` or to `set_parameters`—
#   must be a scalar (or a 1D array of scalars for sweeps).
# - Spinful blocks are produced by callables that assemble $2\times 2$ matrices or Pauli vectors from
#   those scalars.
# - `set_parameters` / `with_parameters` permanently freeze symbolic terms;
#   solver kwargs apply only for that evaluation.
# 
# With these patterns you can combine strings, callables, and shared parameter names
# confidently across both spinless and spinful tight-binding models.