#!/usr/bin/env python3

import matplotlib as mpl
import matplotlib.pyplot as plt

import scipy.integrate as it
from scipy.special import jv as Bessel

import numpy as np

def Main():

    U2 = 0

    x = np.arange(0, 25, 0.1)
    Ut = np.arange(0, 25, 0.1)

    D = np.array([[Density(j, i, U2) for j in x] for i in Ut])

    for i in [0, 10, 20, 50, 100, 200]:
        plt.plot(x, D[i],  marker='.', markersize=0.5, linewidth=0.2, label=f'{Ut[i]}')

    plt.xlabel('$r$')
    plt.ylabel('$n_r(t)$')
    plt.legend(title='Ut')
    plt.savefig('N.pdf')
    plt.close()

    X, Y = np.meshgrid(x, Ut)

    norm = mpl.colors.TwoSlopeNorm(vmin=-4, vcenter=-2, vmax=0)
    Z = np.log(D)

    plt.pcolormesh(X, Y, Z, cmap='hot', linewidth=0, rasterized=True, norm=norm)
    plt.colorbar(label='$\log n_r(t)$')
    plt.xlabel('r')
    plt.ylabel('Ut')
    plt.savefig('Nmap.pdf', dpi=600)

    return

def Density(l, Ut, U2):

    if U2 == 0:
        return Bessel(l, Ut)**2

    ...

    re = ...
    im = ...
    return (re)**2+(im)**2

Main()