N14_Runner.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "019513fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import glob\n",
    "import time\n",
    "import yaml\n",
    "import math\n",
    "import torch as t\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from torch import nn\n",
    "import multiprocessing\n",
    "from torch.optim import Adam\n",
    "from astropy import units as u\n",
    "from astropy import constants as c\n",
    "from torch.autograd import Variable\n",
    "from scipy.interpolate import interp1d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "187ca550",
   "metadata": {},
   "outputs": [],
   "source": [
    "def MP(taskid_lst=None, func=None, Nprocs=24):\n",
    "    def worker(taskid_lst, out_q):\n",
    "        outdict={}\n",
    "        for tid in taskid_lst:\n",
    "            outdict[tid]=func(tid)\n",
    "        out_q.put(outdict)\n",
    "    out_q=multiprocessing.Queue()\n",
    "    chunksize=int(math.ceil(len(taskid_lst)/float(Nprocs)))\n",
    "    procs=[]\n",
    "    for i in range(Nprocs):\n",
    "        p=multiprocessing.Process(target=worker,\\\n",
    "        args=(taskid_lst[chunksize*i:chunksize*(i+1)],out_q))\n",
    "        procs.append(p)\n",
    "        p.start()\n",
    "    resultdict = {}\n",
    "    for i in range(Nprocs):\n",
    "        resultdict.update(out_q.get())\n",
    "    for p in procs:\n",
    "        p.join()\n",
    "    return resultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e1269518",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: UnitsWarning: 'erg/cm^2/K^4/s' contains multiple slashes, which is discouraged by the FITS standard [astropy.units.format.generic]\n"
     ]
    }
   ],
   "source": [
    "sigmaT=c.sigma_T.to('cm^2').value #thompson cross section in cm^2\n",
    "cli=c.c.to('cm/s').value #speed of light in meter/second\n",
    "kbc=c.k_B.to('erg/K').value\n",
    "hli=c.h.to('erg*s').value\n",
    "mele=c.m_e.to('g').value\n",
    "stfBlz=c.sigma_sb.to('erg/cm^2/K^4/s').value #stefan boltzmann constant. \n",
    "secondsInADay=(1*u.day).to('s').value\n",
    "v2d07e16=1/2*(hli**2/2/np.pi/kbc/mele)**(3/2)\n",
    "oneEv=(1*u.eV).to('erg').value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e403b536",
   "metadata": {},
   "outputs": [],
   "source": [
    "yamlName='N14_Model/runner.yml'\n",
    "with open(yamlName,'r') as reader:\n",
    "    howToRun=yaml.safe_load(reader)\n",
    "ModelSaver=howToRun['files']['save_dir']\n",
    "linesFile=howToRun['files']['line_file']\n",
    "levelFile=howToRun['files']['level_file']\n",
    "lineWide=float(howToRun['material']['lines']['line_wide'])\n",
    "veloMin=float(howToRun['material']['velocity']['min'])*cli\n",
    "veloMax=float(howToRun['material']['velocity']['max'])*cli\n",
    "expTime=float(howToRun['material']['density']['exp_time'])*secondsInADay\n",
    "densDataType=howToRun['material']['density']['data_type']\n",
    "if densDataType=='file':\n",
    "    densT0=float(howToRun['material']['density']['t_0'])\n",
    "    densFile=howToRun['material']['density']['file_name']\n",
    "    densData=np.loadtxt(densFile)\n",
    "    veloData=densData[:,0]*100000#Convert from km/s to cm/s. \n",
    "    densData=densData[:,1]*(densT0/expTime)**3\n",
    "    if veloMax>=veloData.max():\n",
    "        print('Gesa Warning: The upper boundary is too high. ')\n",
    "        veloMax=veloData.max()\n",
    "    if veloMin<=veloData.min():\n",
    "        print('Gesa Error: The lower boundary is too low. ')\n",
    "        assert False\n",
    "    radiData=veloData*expTime\n",
    "elif densDataType=='power_law':\n",
    "    dSlope=float(howToRun['material']['density']['slope'])\n",
    "    N0=float(howToRun['material']['density']['pivot_dens'])\n",
    "    V0=float(howToRun['material']['density']['pivot_velo'])*cli\n",
    "    R0=V0*expTime\n",
    "elemDataType=howToRun['material']['element']['data_type']\n",
    "if elemDataType=='file':\n",
    "    elemFile=howToRun['material']['element']['file_name']\n",
    "    elemData=np.loadtxt(elemFile)\n",
    "    veloForElem=elemData[:,0]*100000\n",
    "    radiForElem=veloForElem*expTime\n",
    "    elemData=elemData[:,1:]\n",
    "    maxElem=elemData.shape[1]\n",
    "    elemSum=elemData.sum(axis=0)\n",
    "elif elemDataType=='one_zone':\n",
    "    elemList=t.tensor(np.array(howToRun['material']['element']['number_ratio'])).cuda()\n",
    "    maxElem=len(elemList)\n",
    "    elemSum=elemList\n",
    "extraEngType=howToRun['material']['extra_energy']['data_type']\n",
    "if extraEngType=='file':\n",
    "    extraBounFile=howToRun['material']['extra_energy']['boun_file']\n",
    "    extraSourFile=howToRun['material']['extra_energy']['source_file']\n",
    "    extraBounVal=howToRun['material']['extra_energy']['boun_value']\n",
    "    extraSourVal=howToRun['material']['extra_energy']['source_value']\n",
    "    exBounData=np.loadtxt(extraBounFile)\n",
    "    exSourData=np.loadtxt(extraSourFile)\n",
    "    exEnVelo=exBounData[:,0]*100000\n",
    "    exEnRadi=exEnVelo*expTime\n",
    "    exEnData=exBounData[:,1]*extraBounVal+exSourData[:,1]*extraSourVal\n",
    "    exEnData[exEnData<0]=0\n",
    "elif extraEngType=='zero':\n",
    "    print('Gesa Info: There is no extra energy other than the boundary condition. ')\n",
    "lineLimit=float(howToRun['material']['lines']['line_limit'])\n",
    "allowShrink=howToRun['material']['lines']['allow_shrink']\n",
    "BBBoun=float(howToRun['boundary']['T_inner'])\n",
    "freqMinLog=float(howToRun['boundary']['freq_min'])\n",
    "freqMaxLog=float(howToRun['boundary']['freq_max'])\n",
    "freqSampler=int(howToRun['boundary']['freq_sample'])\n",
    "learnRateSchedule=np.array(howToRun['train']['rate_list'])\n",
    "pdeWeight=float(howToRun['train']['pde_weight'])\n",
    "bo1Weight=float(howToRun['train']['boun1_weight'])\n",
    "bo2Weight=float(howToRun['train']['boun2_weight'])\n",
    "tinyEpochNum=int(howToRun['train']['tiny_epoch'])\n",
    "tempMultiRate=float(howToRun['train']['temp_rate'])\n",
    "l1Norm=float(howToRun['train']['l1_norm'])\n",
    "goodIterate=howToRun['train']['good_iterate']\n",
    "zmin=expTime*veloMin\n",
    "zmax=expTime*veloMax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bccaccbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.exists(ModelSaver):\n",
    "    os.popen('rm -r '+ModelSaver)\n",
    "    time.sleep(0.1)\n",
    "    print('Gesa Warning: Wiping out an old PINN model. ')\n",
    "os.mkdir(ModelSaver)\n",
    "os.mkdir(ModelSaver+'StepModels/')\n",
    "os.mkdir(ModelSaver+'Metrics/')\n",
    "os.popen('cp '+yamlName+' '+ModelSaver+'HyperData.yml')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c34d974",
   "metadata": {},
   "outputs": [],
   "source": [
    "def BBSpec(freqHz,BBBoun=BBBoun):\n",
    "    return 2*hli*freqHz**3/cli**2/(np.e**(hli*freqHz/kbc/BBBoun)-1)\n",
    "freqGrid=np.linspace(freqMinLog,freqMaxLog,num=freqSampler)\n",
    "freqGrid=10**freqGrid\n",
    "freqGrid=t.tensor(freqGrid).cuda()\n",
    "freqWidth=freqGrid*np.log(10)*(freqMaxLog-freqMinLog)/freqSampler\n",
    "freqGridCpu=freqGrid.cpu()\n",
    "freqWidthCpu=freqWidth.cpu()\n",
    "intenMax=BBSpec(freqGrid).max()\n",
    "tempMax=BBBoun\n",
    "threeDimStep=200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41c3feb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "levelData={}\n",
    "elemMask=[]\n",
    "for i in range(maxElem):\n",
    "    elemNum=i+1\n",
    "    levelDataElem={}\n",
    "    ionLimEngs=[]\n",
    "    ionDataMask=[]\n",
    "    if len(glob.glob(levelFile+str(elemNum)+'_'+'*'))==0:\n",
    "        print('Gesa Warning: Element '+str(elemNum)+' has no data. ')\n",
    "        elemMask.append(False)\n",
    "        continue\n",
    "    elif elemSum[i]==0:\n",
    "        print('Gesa Info: Element '+str(elemNum)+' not in the material, not reading the element data. ')\n",
    "        elemMask.append(False)\n",
    "        continue\n",
    "    else:elemMask.append(True)\n",
    "    for j in range(elemNum):\n",
    "        ionNum=j+1\n",
    "        shrinkFile=levelFile+str(elemNum)+'_'+str(ionNum)+'_shrink.csv'\n",
    "        normFile=levelFile+str(elemNum)+'_'+str(ionNum)+'.csv'\n",
    "        if os.path.exists(normFile)==False and os.path.exists(shrinkFile)==False:\n",
    "            print('Gesa Warning: Element '+str(elemNum)+' Ion '+str(ionNum)+' has no level data. ')\n",
    "            ionDataMask.append(False)\n",
    "            continue\n",
    "        elif os.path.exists(normFile) and os.path.exists(shrinkFile)==False:fileName=normFile\n",
    "        elif os.path.exists(shrinkFile) and os.path.exists(normFile)==False:fileName=shrinkFile\n",
    "        elif os.path.exists(shrinkFile) and os.path.exists(normFile):\n",
    "            if allowShrink:fileName=shrinkFile\n",
    "            else:fileName=normFile\n",
    "        ionDataMask.append(True)\n",
    "        tableIon=pd.read_csv(fileName,index_col=0)\n",
    "        limEng=tableIon[tableIon['Configuration']=='Limit']['Level (eV)'].iloc[0]*oneEv\n",
    "        tableIon=tableIon[tableIon['Configuration']!='Limit']\n",
    "        levelDataIon={}\n",
    "        levelDataIon['Config']=np.array(tableIon['Configuration'])\n",
    "        levelDataIon['Term']=np.array(tableIon['Term'].astype('str'))\n",
    "        if tableIon['J'].dtype==float:\n",
    "            if tableIon['J'].isna().mean()==1:levelDataIon['J']=np.array(tableIon['J'].astype('str'))\n",
    "            else:levelDataIon['J']=np.array(tableIon['J'].astype('int').astype('str'))\n",
    "        else:levelDataIon['J']=np.array(tableIon['J'].astype('str'))\n",
    "        levelDataIon['g']=np.array(tableIon['g'])\n",
    "        levelDataIon['Level']=np.array(tableIon['Level (eV)'])*oneEv\n",
    "        levelDataIon['nQua']=np.array(tableIon['nQua'].astype('int'))\n",
    "        ionLimEngs.append(limEng)\n",
    "        levelDataElem[ionNum-1]=levelDataIon\n",
    "    levelDataElem['Limit']=ionLimEngs\n",
    "    levelDataElem['DataMask']=ionDataMask\n",
    "    levelData[elemNum]=levelDataElem\n",
    "levelData['ElemMask']=elemMask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6217e059",
   "metadata": {},
   "outputs": [],
   "source": [
    "lineData={}\n",
    "for i in range(maxElem):\n",
    "    if levelData['ElemMask'][i]==False:continue\n",
    "    elemNum=i+1\n",
    "    lineDataElem={}\n",
    "    for j in range(elemNum):\n",
    "        ionNum=j+1\n",
    "        shrinkFile=linesFile+str(elemNum)+'_'+str(ionNum)+'_shrink.csv'\n",
    "        normFile=linesFile+str(elemNum)+'_'+str(ionNum)+'.csv'\n",
    "        if os.path.exists(normFile)==False and os.path.exists(shrinkFile)==False:\n",
    "            print('Gesa Warning: Element '+str(elemNum)+' Ion '+str(ionNum)+' has no line data. ')\n",
    "            continue\n",
    "        elif os.path.exists(normFile) and os.path.exists(shrinkFile)==False:fileName=normFile\n",
    "        elif os.path.exists(shrinkFile) and os.path.exists(normFile)==False:fileName=shrinkFile\n",
    "        elif os.path.exists(shrinkFile) and os.path.exists(normFile):\n",
    "            if allowShrink:fileName=shrinkFile\n",
    "            else:fileName=normFile\n",
    "        tableIon=pd.read_csv(fileName,index_col=0)\n",
    "        mask=(tableIon['ritz_wl_vac(nm)']<1000)&(tableIon['ritz_wl_vac(nm)']>100)\n",
    "        tableIon=tableIon[mask].reset_index(drop=True)\n",
    "        if len(tableIon)==0:continue\n",
    "        freqIon=(c.c/(np.array(tableIon['ritz_wl_vac(nm)'])*u.nm)).to('Hz').value\n",
    "        mask=(freqIon<10**freqMaxLog)&(freqIon>10**freqMinLog)\n",
    "        tableIon=tableIon[mask]\n",
    "        if len(tableIon)==0:\n",
    "            print('Gesa Warning: Element',elemNum,'Ion',ionNum,' has no line in the wavelength region. ')\n",
    "            continue\n",
    "        lineDataIon={}\n",
    "        lineDataIon['Freq']=(c.c/(np.array(tableIon['ritz_wl_vac(nm)'])*u.nm)).to('Hz').value\n",
    "        lineDataIon['Aki']=np.array(tableIon['Aki(s^-1)'])\n",
    "        lineDataIon['Conf_l']=np.array(tableIon['conf_i'].astype('str'))\n",
    "        lineDataIon['Term_l']=np.array(tableIon['term_i'].astype('str'))\n",
    "        if tableIon['J_i'].dtype==float and np.sum(np.isnan(tableIon['J_i']))==0:lineDataIon['J_l']=np.array(tableIon['J_i'].astype('int').astype('str'))\n",
    "        else:lineDataIon['J_l']=np.array(tableIon['J_i'].astype('str'))\n",
    "        lineDataIon['Conf_u']=np.array(tableIon['conf_k'].astype('str'))\n",
    "        lineDataIon['Term_u']=np.array(tableIon['term_k'].astype('str'))\n",
    "        if tableIon['J_k'].dtype==float and np.sum(np.isnan(tableIon['J_k']))==0:lineDataIon['J_u']=np.array(tableIon['J_k'].astype('int').astype('str'))\n",
    "        else:lineDataIon['J_u']=np.array(tableIon['J_k'].astype('str'))\n",
    "        lineDataIon['Bul']=lineDataIon['Aki']/(2*hli*lineDataIon['Freq']**3/cli**2)\n",
    "        lineDataIon['Blu']=lineDataIon['Bul']*np.array(tableIon['g_k'])/np.array(tableIon['g_i'])\n",
    "        lineDataIon['Mask']=np.ones(len(lineDataIon['Freq']),dtype=bool)\n",
    "        lineDataElem[ionNum-1]=lineDataIon\n",
    "    lineData[elemNum]=lineDataElem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4734204",
   "metadata": {},
   "outputs": [],
   "source": [
    "collBatch=1500\n",
    "bounBatch=1500\n",
    "sourSubBa=1500\n",
    "tempBatch=200\n",
    "batchNum=400\n",
    "epochNum=8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a86b46ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "class IntenW(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(IntenW, self).__init__()\n",
    "        self.fc1=nn.Linear(2,256)\n",
    "        self.fc2=nn.Linear(256,256)\n",
    "        self.fc3=nn.Linear(256,256)\n",
    "        self.fc4=nn.Linear(256,256)\n",
    "        self.fc5=nn.Linear(256,256)\n",
    "        self.fc6=nn.Linear(256,512)\n",
    "        self.fc7=nn.Linear(512,512)\n",
    "        self.fc8=nn.Linear(512,512)\n",
    "        self.fc9=nn.Linear(512,512)\n",
    "        self.fc10=nn.Linear(512,512)\n",
    "        self.fc11=nn.Linear(512,2048)\n",
    "        self.fc12=nn.Linear(2048,2048)\n",
    "        self.fc13=nn.Linear(2048,2048)\n",
    "        self.fc14=nn.Linear(2048,freqSampler)\n",
    "    def forward(self,Xin):\n",
    "        self.Xin=Xin\n",
    "        self.out1=self.fc1(Xin)\n",
    "        self.out2=self.fc2(nn.Tanh()(self.out1))\n",
    "        self.out3=self.fc3(nn.Tanh()(self.out2))\n",
    "        self.out4=self.fc4(nn.Tanh()(self.out3))\n",
    "        self.out5=self.fc5(nn.Tanh()(self.out4))\n",
    "        self.out6=self.fc6(nn.Tanh()(self.out5))\n",
    "        self.out7=self.fc7(nn.Tanh()(self.out6))\n",
    "        self.out8=self.fc8(nn.Tanh()(self.out7))\n",
    "        self.out9=self.fc9(nn.Tanh()(self.out8))\n",
    "        self.out10=self.fc10(nn.Tanh()(self.out9))\n",
    "        self.out11=self.fc11(nn.Tanh()(self.out10))\n",
    "        self.out12=self.fc12(nn.Tanh()(self.out11))\n",
    "        self.out13=self.fc13(nn.Tanh()(self.out12))\n",
    "        self.out14=self.fc14(nn.Tanh()(self.out13))\n",
    "        return self.out14\n",
    "    def partialInput(self):\n",
    "        pout=t.cosh(t.unsqueeze(self.out1,2))**-2*(self.fc1.weight)\n",
    "        pout=t.cosh(t.unsqueeze(self.out2,2))**-2*(self.fc2.weight@pout)\n",
    "        pout=t.cosh(t.unsqueeze(self.out3,2))**-2*(self.fc3.weight@pout)\n",
    "        pout=t.cosh(t.unsqueeze(self.out4,2))**-2*(self.fc4.weight@pout)\n",
    "        pout=t.cosh(t.unsqueeze(self.out5,2))**-2*(self.fc5.weight@pout)\n",
    "        pout=t.cosh(t.unsqueeze(self.out6,2))**-2*(self.fc6.weight@pout)\n",
    "        pout=t.cosh(t.unsqueeze(self.out7,2))**-2*(self.fc7.weight@pout)\n",
    "        pout=t.cosh(t.unsqueeze(self.out8,2))**-2*(self.fc8.weight@pout)\n",
    "        pout=t.cosh(t.unsqueeze(self.out9,2))**-2*(self.fc9.weight@pout)\n",
    "        pout=t.cosh(t.unsqueeze(self.out10,2))**-2*(self.fc10.weight@pout)\n",
    "        pout=t.cosh(t.unsqueeze(self.out11,2))**-2*(self.fc11.weight@pout)\n",
    "        pout=t.cosh(t.unsqueeze(self.out12,2))**-2*(self.fc12.weight@pout)\n",
    "        pout=t.cosh(t.unsqueeze(self.out13,2))**-2*(self.fc13.weight@pout)\n",
    "        pout=self.fc14.weight@pout\n",
    "        return pout\n",
    "class Tempe(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Tempe,self).__init__()\n",
    "        self.fc1=nn.Linear(1,64)\n",
    "        self.fc2=nn.Linear(64,64)\n",
    "        self.fc3=nn.Linear(64,64)\n",
    "        self.fc4=nn.Linear(64,64)\n",
    "        self.fc5=nn.Linear(64,1)\n",
    "    def forward(self,zin):\n",
    "        out=nn.SELU()(self.fc1(zin))\n",
    "        out=nn.SELU()(self.fc2(out))\n",
    "        out=nn.SELU()(self.fc3(out))\n",
    "        out=nn.SELU()(self.fc4(out))\n",
    "        out=-nn.LogSigmoid()(self.fc5(out))\n",
    "        return out\n",
    "TempNet=Tempe().double().cuda()\n",
    "IntenWNetFo=IntenW().double().cuda()\n",
    "IntenWNetFoRec=IntenW().double().cuda()\n",
    "IntenWNetFoRec.load_state_dict(IntenWNetFo.state_dict())\n",
    "IntenWNetBa=IntenW().double().cuda()\n",
    "IntenWNetBaRec=IntenW().double().cuda()\n",
    "IntenWNetBaRec.load_state_dict(IntenWNetBa.state_dict())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78327060",
   "metadata": {},
   "outputs": [],
   "source": [
    "def integGridAssigner(randomGridNumber):\n",
    "    phiGrid=t.tensor(np.random.random(randomGridNumber)*np.pi).cuda()\n",
    "    return phiGrid\n",
    "def veloMaker(zList):\n",
    "    veloList=(veloMax-veloMin)*zList+veloMin\n",
    "    return veloList\n",
    "def DopplerFreqChanger(VeloBeta,phiAngleIn):\n",
    "    GammaRela=1/(1-VeloBeta**2)**0.5\n",
    "    Gamma=GammaRela.reshape([-1,1])\n",
    "    Beta=VeloBeta.reshape([-1,1])\n",
    "    phiAngle=phiAngleIn.reshape([-1,1])\n",
    "    nuBarDivNu=Gamma*(1-t.cos(phiAngle)*Beta)\n",
    "    return nuBarDivNu\n",
    "if densDataType=='file':\n",
    "    densCalc=interp1d(radiData,densData)\n",
    "elif densDataType=='power_law':\n",
    "    def densCalc(radiList):\n",
    "        densList=N0*(radiList/R0)**dSlope\n",
    "        return densList\n",
    "if elemDataType=='file':\n",
    "    elemInterDict={}\n",
    "    for i in range(1,maxElem+1):\n",
    "        elemInterDict[i]=interp1d(radiForElem,elemData[:,i-1])\n",
    "    def elemInterp(radiList,elemNumbers):\n",
    "        elemReturnMat=np.zeros([len(radiList),maxElem])\n",
    "        for i in range(len(elemNumbers)):\n",
    "            elemReturnMat[:,elemNumbers[i]-1]=elemInterDict[elemNumbers[i]](radiList)\n",
    "        return elemReturnMat.T\n",
    "    def elemDensCalc(zList):\n",
    "        radiList=zList.clone().detach().cpu().numpy()*(zmax-zmin)+zmin\n",
    "        densList=densCalc(radiList)\n",
    "        elemList=elemInterp(radiList,np.arange(1,maxElem+1))*densList\n",
    "        return t.tensor(elemList).cuda().T\n",
    "elif elemDataType=='one_zone':\n",
    "    def elemDensCalc(zList):\n",
    "        elemRatioList=t.ones([len(zList),len(elemList)]).cuda()*elemList\n",
    "        radiList=(zList*(zmax-zmin)+zmin).clone().detach().cpu().numpy()\n",
    "        elemDensList=t.tensor(densCalc(radiList).reshape(-1,1)).cuda()*elemRatioList\n",
    "        return elemDensList\n",
    "if extraEngType=='file':\n",
    "    exEnCalc=interp1d(exEnRadi,exEnData)\n",
    "elif extraEngType=='zero':\n",
    "    def exEnCalc(radi):return 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7742d70a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def NiiCalcul(TEmiList,NiList,Ne):\n",
    "    NiiList=[]\n",
    "    for elem in range(1,maxElem+1):\n",
    "        if levelData['ElemMask'][elem-1]==False:\n",
    "            NiiList.append(t.zeros(len(TEmiList)).cuda())\n",
    "            continue\n",
    "        sumHere2=0\n",
    "        for ionJ in range(0,elem+1):\n",
    "            if ionJ<elem:\n",
    "                if levelData[elem]['DataMask'][ionJ]==False:continue\n",
    "            prodHere=1\n",
    "            for ionL in range(ionJ,elem):\n",
    "                if levelData[elem]['DataMask'][ionL]==False:continue\n",
    "                if elem>ionL+1 and levelData[elem]['DataMask'][ionL+1]:gUpper=levelData[elem][ionL+1]['g'][0]\n",
    "                else:gUpper=1\n",
    "                gLower=t.tensor(levelData[elem][ionL]['g']).cuda()\n",
    "                engLevel=t.tensor(levelData[elem]['Limit'][ionL]-levelData[elem][ionL]['Level']).cuda()\n",
    "                sumHere=t.sum(v2d07e16/gUpper*gLower*TEmiList**(-1.5)*np.e**(engLevel/kbc/TEmiList),axis=1)\n",
    "                prodHere=(prodHere*sumHere).detach()*Ne\n",
    "            prodHere=prodHere\n",
    "            sumHere2=sumHere2+prodHere\n",
    "        Nii=NiList[:,elem-1]/sumHere2\n",
    "        NiiList.append(Nii)\n",
    "    return NiiList\n",
    "def NeCalcul(TEmiList,NiiList,neInput):\n",
    "    Ne=0\n",
    "    for elem in range(1,maxElem+1):\n",
    "        if levelData['ElemMask'][elem-1]==False:continue\n",
    "        sumHere2=0\n",
    "        Nii=NiiList[elem-1]\n",
    "        for ionJ in range(0,elem+1):\n",
    "            if ionJ<elem:\n",
    "                if levelData[elem]['DataMask'][ionJ]==False:continue\n",
    "            prodHere=1\n",
    "            for ionL in range(ionJ,elem):\n",
    "                if levelData[elem]['DataMask'][ionL]==False:continue\n",
    "                if elem>ionL+1 and levelData[elem]['DataMask'][ionL+1]:gUpper=levelData[elem][ionL+1]['g'][0]\n",
    "                else:gUpper=1\n",
    "                gLower=t.tensor(levelData[elem][ionL]['g']).cuda()\n",
    "                engLevel=t.tensor(levelData[elem]['Limit'][ionL]-levelData[elem][ionL]['Level']).cuda()\n",
    "                sumHere=t.sum(v2d07e16/gUpper*gLower*TEmiList**(-1.5)*np.e**(engLevel/kbc/TEmiList),axis=1)\n",
    "                prodHere=(prodHere*sumHere).detach()*neInput\n",
    "            prodHere=prodHere*ionJ\n",
    "            sumHere2=sumHere2+prodHere\n",
    "        Ne=Ne+sumHere2*Nii\n",
    "    return Ne\n",
    "def populCalcul(TEmiList,NiList,iterNum=30,initDens=1e4,cudaOutput=True):\n",
    "    if initDens=='1-ion':\n",
    "        neOld=NiList.sum(axis=1).clone().detach()\n",
    "        neNew=NiList.sum(axis=1).clone().detach()\n",
    "    else:\n",
    "        neOld=initDens*t.ones(len(TEmiList)).cuda()\n",
    "        neNew=initDens*t.ones(len(TEmiList)).cuda()\n",
    "    for neRound in range(iterNum):\n",
    "        NiiList=NiiCalcul(TEmiList,NiList,neNew)\n",
    "        Ne=NeCalcul(TEmiList,NiiList,neNew)\n",
    "        neOld=neNew.detach()\n",
    "        neNew=((Ne+neOld)/2).detach()\n",
    "    NiiList=NiiCalcul(TEmiList,NiList,neNew)\n",
    "    neNew=neNew.reshape(-1,1)\n",
    "    popuList={}\n",
    "    for elem in range(1,maxElem+1):\n",
    "        if levelData['ElemMask'][elem-1]==False:continue\n",
    "        Nii=NiiList[elem-1]\n",
    "        popuListIon={}\n",
    "        for ionJ in range(elem,-1,-1):\n",
    "            highlyIonized=np.where(levelData[elem]['DataMask'])[0].max()+1\n",
    "            if ionJ==highlyIonized:popuListIon[ionJ]=NiiList[elem-1].reshape(-1,1)\n",
    "            elif ionJ>highlyIonized:continue\n",
    "            else:\n",
    "                if ionJ+1==elem:gUpper=1\n",
    "                elif levelData[elem]['DataMask'][ionJ+1]==False:gUpper=1\n",
    "                else:gUpper=levelData[elem][ionJ+1]['g'][0]\n",
    "                gLower=t.tensor(levelData[elem][ionJ]['g']).cuda()\n",
    "                engLevel=t.tensor(levelData[elem]['Limit'][ionJ]-levelData[elem][ionJ]['Level']).cuda()\n",
    "                popuListIon[ionJ]=popuListIon[ionJ+1][:,[0]]*neNew*v2d07e16/gUpper*gLower*TEmiList**(-1.5)*np.e**(engLevel/kbc/TEmiList)\n",
    "        popuList[elem]=popuListIon\n",
    "    if cudaOutput==False:\n",
    "        neNew=neNew.detach().cpu().numpy()\n",
    "        for elem in popuList.keys():\n",
    "            for ionJ in popuList[elem].keys():\n",
    "                popuList[elem][ionJ]=popuList[elem][ionJ].detach().cpu().numpy()\n",
    "    return neNew,popuList"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d28901bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def BBsourceMakerSparse(lineWide,nuBarDivNu,popuDens):\n",
    "    freqList=freqGrid*nuBarDivNu\n",
    "    jEmiLine=0\n",
    "    kAbsLine=0\n",
    "    for elem in lineData.keys():\n",
    "        if levelData['ElemMask'][elem-1]==False:continue\n",
    "        for ionJ in lineData[elem].keys():\n",
    "            if levelData[elem]['DataMask'][ionJ]==False:continue\n",
    "            lineMask=lineData[elem][ionJ]['Mask']\n",
    "            if np.sum(lineMask)==0:continue\n",
    "            lineFreq=t.tensor(lineData[elem][ionJ]['Freq'][lineMask]).cuda().reshape(-1,1,1)\n",
    "            Aul=t.tensor(lineData[elem][ionJ]['Aki'][lineMask]).cuda().reshape(-1,1,1)\n",
    "            Blu=t.tensor(lineData[elem][ionJ]['Blu'][lineMask]).cuda().reshape(-1,1,1)\n",
    "            Bul=t.tensor(lineData[elem][ionJ]['Bul'][lineMask]).cuda().reshape(-1,1,1)\n",
    "            popuDensOne=popuDens[elem][ionJ]\n",
    "            levelFormat=levelData[elem][ionJ]['Config']+levelData[elem][ionJ]['Term']+levelData[elem][ionJ]['J']\n",
    "            levelUpper=lineData[elem][ionJ]['Conf_u'][lineMask]+lineData[elem][ionJ]['Term_u'][lineMask]+lineData[elem][ionJ]['J_u'][lineMask]\n",
    "            levelLower=lineData[elem][ionJ]['Conf_l'][lineMask]+lineData[elem][ionJ]['Term_l'][lineMask]+lineData[elem][ionJ]['J_l'][lineMask]\n",
    "            upperMask=np.where(levelUpper.reshape(-1,1)==levelFormat)[1]\n",
    "            lowerMask=np.where(levelLower.reshape(-1,1)==levelFormat)[1]\n",
    "            maskHere=(t.abs(freqList-lineFreq)<=freqWidth*lineWide)\n",
    "            if t.sum(maskHere)==0:continue\n",
    "            maskHere=maskHere.to_sparse()\n",
    "            gausLineSha=t.sparse.FloatTensor(maskHere._indices(),maskHere._values()/(lineWide*2)/freqWidth[maskHere._indices()[2]],maskHere.size())\n",
    "            gaInd=gausLineSha._indices()\n",
    "            gaVal=gausLineSha._values()\n",
    "            jEmiLineOne=t.sparse.FloatTensor(gaInd,gaVal*Aul[gaInd[0],0,0]*lineFreq[gaInd[0],0,0]*popuDensOne[:,upperMask].T[gaInd[0],gaInd[1]],gausLineSha.size())\n",
    "            try:jEmiLineOne=t.sparse.sum(jEmiLineOne*hli/4/np.pi,0).to_dense()/nuBarDivNu**2\n",
    "            except:print(jEmiLineOne,maskHere,freqList,lineFreq)\n",
    "            jEmiLine=jEmiLine+jEmiLineOne.detach().cpu().numpy()\n",
    "            del jEmiLineOne\n",
    "            kAbsLineOne=(popuDensOne[:,lowerMask].T.unsqueeze(2)*Blu-popuDensOne[:,upperMask].T.unsqueeze(2)*Bul)\n",
    "            kAbsLineOne=t.sparse.FloatTensor(gaInd,gaVal*lineFreq[gaInd[0],0,0]*kAbsLineOne[gaInd[0],gaInd[1],0],gausLineSha.size())\n",
    "            kAbsLineOne=t.sparse.sum(kAbsLineOne*hli/4/np.pi,0).to_dense()*nuBarDivNu\n",
    "            kAbsLine=kAbsLine+kAbsLineOne.detach().cpu().numpy()\n",
    "            del maskHere\n",
    "            del gausLineSha\n",
    "            del kAbsLineOne\n",
    "            del gaInd\n",
    "            del gaVal\n",
    "    return jEmiLine,kAbsLine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb2fe377",
   "metadata": {},
   "outputs": [],
   "source": [
    "def makeAllTheSource(zftp,TempNetIn):\n",
    "    TEmiList=TempNetIn(zftp[:,[0]]).detach()*tempMax\n",
    "    VeloBeta=veloMaker(zftp[:,0])/cli\n",
    "    nuBarDivNu=DopplerFreqChanger(VeloBeta,zftp[:,1])\n",
    "    NiList=elemDensCalc(zftp[:,0])\n",
    "    elecDens,popuDens=populCalcul(TEmiList,NiList,initDens='1-ion')\n",
    "    sigList=elecDens*sigmaT\n",
    "    thermSourc=sigList.reshape([-1,1])*BBSpec(freqGrid*nuBarDivNu,TEmiList.reshape([-1,1]))\n",
    "    thermSourc=thermSourc.clone().detach()\n",
    "    thermSourc[t.isnan(thermSourc)]=0\n",
    "    thermSourc=thermSourc/nuBarDivNu**2\n",
    "    jEmiLine,kAbsLine=BBsourceMakerSparse(lineWide,nuBarDivNu,popuDens)\n",
    "    jEmiLine=t.tensor(jEmiLine).cuda()\n",
    "    kAbsLine=t.tensor(kAbsLine).cuda()\n",
    "    totalAbs=sigList+kAbsLine\n",
    "    totalEmi=thermSourc+jEmiLine\n",
    "    return totalAbs,totalEmi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39c53b69",
   "metadata": {},
   "outputs": [],
   "source": [
    "def subBatchMakeSource(zftp,TempNetIn):\n",
    "    totalAbs,totalEmi=[],[]\n",
    "    zftpSplit=zftp.split(sourSubBa)\n",
    "    for subBaNum in range(len(zftpSplit)):\n",
    "        subBaAbs,subBaEmi=makeAllTheSource(zftpSplit[subBaNum],TempNetIn)\n",
    "        totalAbs.append(subBaAbs)\n",
    "        totalEmi.append(subBaEmi)\n",
    "    totalAbs=t.cat(totalAbs)\n",
    "    totalEmi=t.cat(totalEmi)\n",
    "    return totalAbs,totalEmi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13e47828",
   "metadata": {},
   "outputs": [],
   "source": [
    "def lineMaskMakeOne(elem,ionJ,popuDensOne,lineLimit):\n",
    "    nuBarDivNu=t.ones([len(popuDensOne),1])\n",
    "    freqList=freqGridCpu*nuBarDivNu\n",
    "    try:lineFreq=t.tensor(lineData[elem][ionJ]['Freq']).reshape(-1,1,1)\n",
    "    except:print(elem,ionJ)\n",
    "    Aul=t.tensor(lineData[elem][ionJ]['Aki']).reshape(-1,1,1)\n",
    "    Blu=t.tensor(lineData[elem][ionJ]['Blu']).reshape(-1,1,1)\n",
    "    Bul=t.tensor(lineData[elem][ionJ]['Bul']).reshape(-1,1,1)\n",
    "    levelFormat=levelData[elem][ionJ]['Config']+levelData[elem][ionJ]['Term']+levelData[elem][ionJ]['J']\n",
    "    levelUpper=lineData[elem][ionJ]['Conf_u']+lineData[elem][ionJ]['Term_u']+lineData[elem][ionJ]['J_u']\n",
    "    levelLower=lineData[elem][ionJ]['Conf_l']+lineData[elem][ionJ]['Term_l']+lineData[elem][ionJ]['J_l']\n",
    "    upperMask=np.where(levelUpper.reshape(-1,1)==levelFormat)[1]\n",
    "    lowerMask=np.where(levelLower.reshape(-1,1)==levelFormat)[1]\n",
    "    maskHere=(t.abs(freqList-lineFreq)<=freqWidthCpu*lineWide).to_sparse()\n",
    "    gausLineSha=t.sparse.FloatTensor(maskHere._indices(),maskHere._values()/(lineWide*2)/freqWidthCpu[maskHere._indices()[2]],maskHere.size())\n",
    "    gaInd=gausLineSha._indices()\n",
    "    gaVal=gausLineSha._values()\n",
    "    kAbsLineOne=(popuDensOne[:,lowerMask].T.unsqueeze(2)*Blu-popuDensOne[:,upperMask].T.unsqueeze(2)*Bul)\n",
    "    kAbsLineOne=t.sparse.FloatTensor(gaInd,gaVal*lineFreq[gaInd[0],0,0]*kAbsLineOne[gaInd[0],gaInd[1],0],gausLineSha.size())\n",
    "    kAbsLineOne=kAbsLineOne.to_dense()*hli/4/np.pi\n",
    "    kAbsLineOne=(kAbsLineOne.sum(axis=1)*(zmax-zmin)/len(popuDensOne)).cpu().numpy()\n",
    "    maskOpac=(kAbsLineOne.sum(axis=0)>lineLimit)\n",
    "    maskTrans=(kAbsLineOne[:,maskOpac].sum(axis=1)>0)\n",
    "    return maskTrans\n",
    "def lineMaskMake(lineDataIn,temperNet,sampleNum=3500):\n",
    "    lineDataOut=lineDataIn.copy()\n",
    "    zftp=np.zeros([sampleNum,2])\n",
    "    zftp[:,0]=np.linspace(0,1,num=sampleNum)\n",
    "    zftp=Variable(t.tensor(zftp).cuda(),requires_grad=True)\n",
    "    TEmiList=temperNet(zftp[:,[0]]).detach()*tempMax\n",
    "    TEmiList=t.tensor(TEmiList.cpu().detach().numpy()).cuda()\n",
    "    NiList=elemDensCalc(zftp[:,0])\n",
    "    elecDens,popuDens=populCalcul(TEmiList,NiList,initDens='1-ion',cudaOutput=False)\n",
    "    TEmiList=TEmiList.cpu()\n",
    "    def mpInRun(inKey):\n",
    "        elem=int(inKey.split('_')[0])\n",
    "        ionJ=int(inKey.split('_')[1])\n",
    "        popuDensOne=t.tensor(popuDens[elem][ionJ])\n",
    "        maskTrans=lineMaskMakeOne(elem,ionJ,popuDensOne,lineLimit)\n",
    "        return maskTrans\n",
    "    inKeyList=[]\n",
    "    for elem in lineDataIn.keys():\n",
    "        if levelData['ElemMask'][elem-1]==False:continue\n",
    "        for ionJ in lineDataIn[elem].keys():\n",
    "            if levelData[elem]['DataMask'][ionJ]==False:continue\n",
    "            inKeyList.append(str(elem)+'_'+str(ionJ))\n",
    "    mpOut=MP(inKeyList,mpInRun)\n",
    "    for inKey in inKeyList:\n",
    "        elem,ionJ=inKey.split('_')\n",
    "        elem,ionJ=int(elem),int(ionJ)\n",
    "        lineDataOut[elem][ionJ]['Mask']=mpOut[str(elem)+'_'+str(ionJ)]\n",
    "    return lineDataOut"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ab4c239",
   "metadata": {},
   "outputs": [],
   "source": [
    "zftpOrig=np.random.random(collBatch*batchNum*2).reshape(-1,2)\n",
    "zftpOrig=zftpOrig*np.array([1,np.pi])+np.array([0,0])\n",
    "zftpOrig=t.tensor(zftpOrig).cuda()\n",
    "zftpOrig=zftpOrig.split(collBatch)\n",
    "zftpOrig=[i for i in zftpOrig]\n",
    "\n",
    "zOrig=np.random.random(tempBatch*batchNum*1).reshape(-1,1)\n",
    "zOrig=zOrig*np.array([1])\n",
    "zOrig=t.tensor(zOrig).cuda()\n",
    "zOrig=zOrig.split(tempBatch)\n",
    "zOrig=[i for i in zOrig]\n",
    "\n",
    "zftpBoun1Orig=np.random.random(bounBatch*batchNum*2).reshape(-1,2)\n",
    "zftpBoun1Orig=zftpBoun1Orig*np.array([0,np.pi*0.5])+np.array([0,0])\n",
    "zftpBoun1Orig=t.tensor(zftpBoun1Orig).cuda()\n",
    "zftpBoun1Orig=zftpBoun1Orig.split(bounBatch)\n",
    "zftpBoun1Orig=[i for i in zftpBoun1Orig]\n",
    "\n",
    "zftpBoun2Orig=np.random.random(bounBatch*batchNum*2).reshape(-1,2)\n",
    "zftpBoun2Orig=zftpBoun2Orig*np.array([0,np.pi*0.5])+np.array([1,np.pi*0.5])\n",
    "zftpBoun2Orig=t.tensor(zftpBoun2Orig).cuda()\n",
    "zftpBoun2Orig=zftpBoun2Orig.split(bounBatch)\n",
    "zftpBoun2Orig=[i for i in zftpBoun2Orig]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6d95bc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "lossPFoAll=[]\n",
    "lossPBaAll=[]\n",
    "lossBoFAll=[]\n",
    "lossBoBAll=[]\n",
    "lossTemAll=[]\n",
    "lossL1FAll=[]\n",
    "lossL1BAll=[]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46524ba2",
   "metadata": {},
   "outputs": [],
   "source": [
    "for bigStep in range(len(learnRateSchedule)):\n",
    "    lineData=lineMaskMake(lineData,TempNet,100)\n",
    "    optimFo=Adam(IntenWNetFo.parameters(),lr=learnRateSchedule[bigStep])\n",
    "    optimBa=Adam(IntenWNetBa.parameters(),lr=learnRateSchedule[bigStep])\n",
    "    optimTem=Adam(TempNet.parameters(),lr=learnRateSchedule[bigStep]*tempMultiRate)#\n",
    "    for epoch in range(epochNum):\n",
    "        for batch in range(batchNum):\n",
    "            zftpTwo=zftpOrig[batch].clone().detach()\n",
    "            zRadiTwo=zftpTwo[:,[0]]*(zmax-zmin)+zmin\n",
    "            touchBoun=((zRadiTwo*t.sin(zftpTwo[:,[1]]))<zmin)&(zftpTwo[:,[1]]<np.pi/2)\n",
    "            zftp=zftpTwo[touchBoun.flatten()]\n",
    "            zRadi=zRadiTwo[touchBoun.flatten()]\n",
    "            totalAbs,totalEmi=subBatchMakeSource(zftp,TempNet)\n",
    "            for tinyStep in range(tinyEpochNum):\n",
    "                IntenWNetFo.zero_grad()\n",
    "                IntenOut=IntenWNetFo(zftp)*intenMax\n",
    "                XinGrad=IntenW.partialInput(IntenWNetFo)*intenMax\n",
    "                XinGrad[:,:,0]=XinGrad[:,:,0]/(zmax-zmin)\n",
    "                lossEqu=t.cos(zftp[:,[1]])*XinGrad[:,:,0]-t.sin(zftp[:,[1]])*XinGrad[:,:,1]/zRadi+totalAbs*IntenOut-totalEmi\n",
    "                lossEqu=(lossEqu/intenMax/totalAbs.mean())**2\n",
    "                lossEqu=t.mean(lossEqu)*pdeWeight\n",
    "                lossPFoAll.append(lossEqu.item())\n",
    "                zftpBoun1=zftpBoun1Orig[batch].clone().detach()\n",
    "                IntenOutBoun1=IntenWNetFo(zftpBoun1)*intenMax\n",
    "                realBoun=t.tensor(BBSpec(freqGrid)).cuda()\n",
    "                lossBoun1=((IntenOutBoun1-realBoun)/intenMax)**2\n",
    "                lossBoun1=t.mean(lossBoun1)*bo1Weight#(10)\n",
    "                lossBoFAll.append(lossBoun1.item())\n",
    "                \n",
    "                l1Loss=0\n",
    "                for parOne in IntenWNetFo.parameters():\n",
    "                    l1Loss=l1Loss+t.mean(t.abs(parOne))\n",
    "                l1Loss=l1Loss*l1Norm\n",
    "                lossL1FAll.append(l1Loss.item())\n",
    "                #lossSym=t.mean((t.abs(IntenOut)-IntenOut)/intenMax)*100000\n",
    "                loss=lossEqu+lossBoun1+l1Loss#+lossSym\n",
    "                assert np.isnan(loss.item())==False\n",
    "                if goodIterate:\n",
    "                    if tinyStep==0:lossGood=loss.item()\n",
    "                    else:\n",
    "                        if loss.item()<lossGood:\n",
    "                            IntenWNetFoRec.load_state_dict(IntenWNetFo.state_dict())\n",
    "                            lossGood=loss.item()\n",
    "                    loss.backward()\n",
    "                    optimFo.step()\n",
    "                else:\n",
    "                    loss.backward()\n",
    "                    optimFo.step()\n",
    "                    if tinyStep==tinyEpochNum-1:\n",
    "                        IntenWNetFoRec.load_state_dict(IntenWNetFo.state_dict())\n",
    "            \n",
    "            zftp=zftpTwo[(touchBoun==False).flatten()]\n",
    "            zRadi=zRadiTwo[(touchBoun==False).flatten()]\n",
    "            totalAbs,totalEmi=subBatchMakeSource(zftp,TempNet)\n",
    "            for tinyStep in range(tinyEpochNum):\n",
    "                IntenWNetBa.zero_grad()\n",
    "                IntenOut=IntenWNetBa(zftp)*intenMax\n",
    "                XinGrad=IntenW.partialInput(IntenWNetBa)*intenMax\n",
    "                XinGrad[:,:,0]=XinGrad[:,:,0]/(zmax-zmin)\n",
    "                lossEqu=t.cos(zftp[:,[1]])*XinGrad[:,:,0]-t.sin(zftp[:,[1]])*XinGrad[:,:,1]/zRadi+totalAbs*IntenOut-totalEmi\n",
    "                lossEqu=(lossEqu/intenMax/totalAbs.mean())**2\n",
    "                lossEqu=t.mean(lossEqu)*pdeWeight\n",
    "                lossPBaAll.append(lossEqu.item())\n",
    "            \n",
    "                zftpBoun2=zftpBoun2Orig[batch].clone().detach()\n",
    "                IntenOutBoun2=IntenWNetBa(zftpBoun2)*intenMax\n",
    "                lossBoun2=((IntenOutBoun2-0)/intenMax)**2\n",
    "                lossBoun2=t.mean(lossBoun2)*bo2Weight\n",
    "                lossBoBAll.append(lossBoun2.item())\n",
    "                \n",
    "                l1Loss=0\n",
    "                for parOne in IntenWNetBa.parameters():\n",
    "                    l1Loss=l1Loss+t.mean(t.abs(parOne))\n",
    "                l1Loss=l1Loss*l1Norm\n",
    "                lossL1BAll.append(l1Loss.item())\n",
    "                #lossSym=t.mean((t.abs(IntenOut)-IntenOut)/intenMax)*100000\n",
    "                loss=lossEqu+lossBoun2+l1Loss#+lossSym\n",
    "                assert np.isnan(loss.item())==False\n",
    "                if goodIterate:\n",
    "                    if tinyStep==0:lossGood=loss.item()\n",
    "                    else:\n",
    "                        if loss.item()<lossGood:\n",
    "                            IntenWNetBaRec.load_state_dict(IntenWNetBa.state_dict())\n",
    "                            lossGood=loss.item()\n",
    "                    loss.backward()\n",
    "                    optimBa.step()\n",
    "                else:\n",
    "                    loss.backward()\n",
    "                    optimBa.step()\n",
    "                    if tinyStep==tinyEpochNum-1:\n",
    "                        IntenWNetBaRec.load_state_dict(IntenWNetBa.state_dict())\n",
    "            \n",
    "            IntenWNetFo.load_state_dict(IntenWNetFoRec.state_dict())\n",
    "            IntenWNetBa.load_state_dict(IntenWNetBaRec.state_dict())\n",
    "            zTemp=zOrig[batch].clone().detach()\n",
    "            radiTemp=zTemp*(zmax-zmin)+zmin\n",
    "            TempNet.zero_grad()\n",
    "            phiGrid=integGridAssigner(threeDimStep)\n",
    "            phiIn,zbig=t.meshgrid(phiGrid,zTemp[:,0])\n",
    "            zftpBig=t.vstack([zbig.flatten(),phiIn.flatten()]).T\n",
    "            zRadiBig=zftpBig[:,[0]]*(zmax-zmin)+zmin\n",
    "            touchBoun=((zRadiBig*t.sin(zftpBig[:,[1]]))<zmin)&(zftpBig[:,[1]]<np.pi/2)\n",
    "\n",
    "            IntenBig=(IntenWNetFo(zftpBig)*touchBoun+IntenWNetBa(zftpBig)*(1-1*touchBoun)).reshape(threeDimStep,-1,freqSampler)*intenMax\n",
    "            IntenBig=IntenBig.clone().detach()\n",
    "            EnergyAbsoSum=t.sum(IntenBig*t.unsqueeze(t.sin(phiIn),2)*freqGrid,axis=(0,2))*np.log(10)*(freqMaxLog-freqMinLog)*2*np.pi**2/threeDimStep/freqSampler\n",
    "            EnergyAbsoSum=nn.ReLU()(EnergyAbsoSum.reshape([-1,1])-0.001)+0.001\n",
    "            addExtraEnergy=t.tensor(exEnCalc(radiTemp.clone().detach().cpu().numpy())).cuda()\n",
    "\n",
    "            TEmiPred=TempNet(zTemp)*tempMax\n",
    "            NiList=elemDensCalc(zTemp[:,0])\n",
    "            elecDens,popuDens=populCalcul(TEmiPred,NiList,initDens='1-ion')\n",
    "            sigList=elecDens*sigmaT\n",
    "            TEmiTrue=((addExtraEnergy/sigList+EnergyAbsoSum)/4/stfBlz)**0.25\n",
    "            TEmiTrue=t.tensor(TEmiTrue.clone().detach().cpu().numpy()).cuda()\n",
    "\n",
    "            lossTem=t.sum((TEmiTrue-TEmiPred)**2)*1e-6\n",
    "            lossTemAll.append(lossTem.item())\n",
    "            lossTem.backward()\n",
    "            optimTem.step()\n",
    "            print(loss.item(),lossEqu.item(),lossBoun1.item(),lossBoun2.item(),lossTem.item(),l1Loss.item())\n",
    "            \n",
    "        print('epoch '+str(epoch))\n",
    "        t.save(IntenWNetFoRec.state_dict(),ModelSaver+'CurrentIntenForward.to')\n",
    "        t.save(IntenWNetBaRec.state_dict(),ModelSaver+'CurrentIntenBackward.to')\n",
    "        t.save(TempNet.state_dict(),ModelSaver+'CurrentTemp.to')\n",
    "        np.save(ModelSaver+'Metrics/lossPFoAll.npy',np.array(lossPFoAll))\n",
    "        np.save(ModelSaver+'Metrics/lossPBaAll.npy',np.array(lossPBaAll))\n",
    "        np.save(ModelSaver+'Metrics/lossBoFAll.npy',np.array(lossBoFAll))\n",
    "        np.save(ModelSaver+'Metrics/lossBoBAll.npy',np.array(lossBoBAll))\n",
    "        np.save(ModelSaver+'Metrics/lossTemAll.npy',np.array(lossTemAll))\n",
    "        np.save(ModelSaver+'Metrics/lossL1FAll.npy',np.array(lossL1FAll))\n",
    "        np.save(ModelSaver+'Metrics/lossL1BAll.npy',np.array(lossL1BAll))\n",
    "    t.save(IntenWNetFoRec.state_dict(),ModelSaver+'StepModels/IntenForward_'+str(bigStep)+'.to')\n",
    "    t.save(IntenWNetBaRec.state_dict(),ModelSaver+'StepModels/IntenBackward_'+str(bigStep)+'.to')\n",
    "    t.save(TempNet.state_dict(),ModelSaver+'StepModels/Temp_'+str(bigStep)+'.to')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "223316fe",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "389d6383",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:.conda-pytorch] *",
   "language": "python",
   "name": "conda-env-.conda-pytorch-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}