<jupyter_start><jupyter_text>We build the network with an input shape and print out a summary.<jupyter_code>tf.keras.backend.clear_session()

unet_instance = hp_unet.HealpyUNet(nside, indices, learning_rate, mask, mean, n_neighbors=20)
unet_model = unet_instance.model()<jupyter_output>Model: "functional_1"
__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_1 (InputLayer)            [(None, 12288, 1)]   0                                            
__________________________________________________________________________________________________
chebyshev (Chebyshev)           (None, 12288, 16)    208         input_1[0][0]                    
__________________________________________________________________________________________________
healpy_pool (HealpyPool)        (None, 3072, 16)     0           chebyshev[0][0]                  
__________________________________________________________________________________________________
chebyshev_1 (Chebyshev)         (None, 3072, 32)     5216        healpy_pool[0][0]     [...]<jupyter_text>We create a dataset for the training with a channel dimension (not strictly necessary if you use the TF routines for training) and labels. We run the model.<jupyter_code>tot_train_loss=[]
tot_val_loss=[]
import time

epoch=60
p=0.95 #train-test split

for k in range(epoch):
    start=time.time()
    
    ang1, ang2, ang3 = np.random.uniform(high=360,size=(3,gaussian_maps.shape[0]))

    y = np.array([rotate(j,ang1[i],ang2[i],ang3[i],nside,p=1) for i,j in enumerate(gaussian_maps)])
    
    x = y + np.random.normal(0,0.04,y.shape)

    x = np.array([np.where(mask>0.5,x_i,mean) for x_i in x])
    y = np.array([np.where(mask>0.5,y_i,mean) for y_i in y])

    x = x.astype(np.float32)[..., None]
    y = y.astype(np.float32)[..., None]
    
    n_train = int(x.shape[0]*p)
    
    print('Epoch {}/{}'.format(k+1,epoch))
    history = unet_model.fit(
        x=x[:n_train],
        y=y[:n_train],
        batch_size=batch_size,
        epochs=1,
        validation_data=(x[n_train:],y[n_train:])        
    )
    tot_train_loss.append(history.history["loss"][0])
    tot_val_loss.append(history.history["val_loss"][0])

    print(time.time()-start)<jupyter_output>Epoch 1/60
 2/95 [..............................] - ETA: 57s - loss: 3.1731e-04WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0080s vs `on_train_batch_end` time: 0.6009s). Check your callbacks.
95/95 [==============================] - ETA: 0s - loss: 3.0022e-04WARNING:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0043s vs `on_test_batch_end` time: 0.3068s). Check your callbacks.
95/95 [==============================] - 59s 625ms/step - loss: 3.0022e-04 - val_loss: 3.2823e-04
77.49900579452515
Epoch 2/60
 2/95 [..............................] - ETA: 56s - loss: 2.9697e-04WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0077s vs `on_train_batch_end` time: 0.6015s). Check your callbacks.
95/95 [==============================] - ETA: 0s - loss: 2.8545e-04WARNING:tensorflow:Callbacks method `on_test_batch_end` is slow compa[...]<jupyter_text>A final evaluation on the test set.<jupyter_code>unet_model.evaluate(x[n_train:],y[n_train:], batch_size)<jupyter_output>2/5 [===========>..................] - ETA: 0s - loss: 1.1872e-04WARNING:tensorflow:Callbacks method `on_test_batch_end` is slow compared to the batch time (batch time: 0.0049s vs `on_test_batch_end` time: 0.3093s). Check your callbacks.
5/5 [==============================] - 2s 314ms/step - loss: 1.1671e-04<jupyter_text>Let's plot the loss.(Note that the history object contains much more info.)<jupyter_code>plt.figure(figsize=(12,8))
plt.plot(tot_train_loss, label="training")
plt.plot(tot_val_loss, label="validation")
plt.grid()
plt.yscale("log")
plt.legend()
plt.xlabel("Epoch")
plt.ylabel("Loss")
prediction = unet_model.predict(x[-1][None,...])

cm = plt.cm.RdBu_r

hp.mollview(x[-1].flatten(), title='Input', nest=True, cmap=cm, min=np.min(y[-1]), max=np.max(y[-1]))
hp.graticule()
hp.mollview(prediction[0].flatten(), title='Prediction', nest=True, cmap=cm, min=np.min(y[-1]), max=np.max(y[-1]))
hp.graticule()
hp.mollview(y[-1].flatten(), title='Real', nest=True, cmap=cm, min=np.min(y[-1]), max=np.max(y[-1]))
hp.graticule()

prediction = np.array([prediction[0].flatten()])
prediction = np.array([np.where(mask>0.5,x_i,mean) for x_i in prediction])

hp.mollview(prediction[0], title='Prediction Masked', nest=True, cmap=cm, min=np.min(y[-1]), max=np.max(y[-1]))
hp.graticule()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.19696048632218846, 'fraction_numerical': 0.05167173252279635, 'mean_word_length': 4.088659793814433, 'pattern_counts': {'":': 0, '<': 14, '<?xml version=': 0, '>': 18, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '146811', 'license': 'MIT', 'n_tokens_mistral': 1828, 'n_tokens_neox': 1649, 'n_words': 331, 'path': 'deepmass/spherical_mapping/unet_model.ipynb', 'repo_name': 'david-w-a/DeepMass', 'size': '4935'}