<jupyter_start><jupyter_text>Modules<jupyter_code># Plot training get_loss_hist
def plot_train_hist(model, val=False, y_scale='linear'):
    x1 = np.arange(1, len(model.get_loss_hist())+1)
    plt.scatter(x1, model.get_loss_hist(), marker='.', label='Training loss' )
    #if val:
    #    plt.scatter(x1, model.get_loss_hist['val_loss'], marker='.', label='Validation loss')
    plt.xlabel('Training epochs')
    plt.ylabel('Loss')
    plt.yscale(y_scale)
    plt.legend()
    print("Final loss: ", model.get_loss_hist()[-1])
    return plt.show()
def cal_total_loss(dataset, kahler_pot, loss_function, c1):
    total_loss = tf.constant(0, dtype=tf.float32)
    total_mass = tf.constant(0, dtype=tf.float32)
    
    for step, (points, Omega_Omegabar, mass, restriction) in enumerate(dataset):
        curv_tr = curvature_trace(points, kahler_pot, hermitian_met, mass)
        mass_sum = tf.reduce_sum(mass)
        total_loss += loss_func(c1, curv_tr, mass)  * mass_sum
        total_mass += mass_sum
    total_loss = total_loss / total_mass

    return total_loss.numpy()<jupyter_output><empty_output><jupyter_text>Networks for $\mathcal{O}_X(1)$, $\mathcal{O}_X(2)$, $\mathcal{O}_X(4)$ respectively.\begin{itemize}\item "dim" is the number of projective coordinates of points on the CY hypersurface.\item "vec" is a vector giving the number of neurons for layers 1 through d-1\item "rank" is the rank of the vector bundle -- so far should just be 1\end{itemize}<jupyter_code>class Hermitian_metric_O1(tf.keras.Model):
    def __init__(self, dim, rank, vec):
        super(Hermitian_metric_O1, self).__init__()
        self.dim = dim
        self.rank = rank
        self.vec = vec
        self.loss_hist = []
        print("Bundle rank", self.rank)
        self.bihomogeneous = bnn.Bihomogeneous(self.dim)
        self.layer1 = bnn.Dense(self.dim**2, self.rank)
    
    def append_loss(self, loss):
        self.loss_hist.append([loss])
    
    def get_loss_hist(self):
        return self.loss_hist
    
    def call(self, inputs):
        x = self.bihomogeneous(inputs)
        x = tf.math.log(x)
        return x
    

class Hermitian_metric_O2(tf.keras.Model):
    def __init__(self, dim, rank, vec):
        super(Hermitian_metric_O2, self).__init__()
        self.dim = dim
        self.rank = rank
        self.vec = vec
        self.loss_hist = []
        print("Bundle rank", self.rank)
        self.bihomogeneous = bnn.Bihomogeneous(self.dim)
        self.layer1 = bnn.Dense(self.dim**2, self.vec[0], activation=tf.square)
        self.layer2 = bnn.Dense(self.vec[0], self.rank)
        
    def append_loss(self, loss):
        self.loss_hist.append([loss])
    
    def get_loss_hist(self):
        return self.loss_hist
    
    def call(self, inputs):
        x = self.bihomogeneous(inputs)
        x = self.layer1(x)
        x = self.layer2(x)
        x = tf.math.log(x)
        return x
  

class Hermitian_metric_O4(tf.keras.Model):
    def __init__(self, dim, rank, vec):
        super(Hermitian_metric_O4, self).__init__()
        self.dim = dim
        self.rank = rank
        self.vec = vec
        self.loss_hist = []
        print("Bundle rank", self.rank)
        self.bihomogeneous = bnn.Bihomogeneous(self.dim)
        self.layer1 = bnn.Dense(dim**2, self.vec[0], activation=tf.square)
        self.layer2 = bnn.Dense(self.vec[0], self.vec[1], activation=tf.square)
        self.layer3 = bnn.Dense(self.vec[1], self.rank)
    
    def append_loss(self, loss):
        self.loss_hist.append([loss])
    
    def get_loss_hist(self):
        return self.loss_hist
    
    def call(self, inputs):
        x = self.bihomogeneous(inputs)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = tf.math.log(x)
        return x<jupyter_output><empty_output><jupyter_text>Curvature calculation \begin{itemize}\item Issue with normalization needs to be fixed here, in weighted_loss\item Other issue is taking inverse of K\"ahler metric or curvature tensor when calculating $g^{i\bar{j}} F_{i\bar{j}}$\end{itemize}<jupyter_code>def weighted_MAPE(y_true, y_pred, mass):
    weights = mass / K.sum(mass)
    return K.sum(K.abs(y_true - y_pred) / tf.cast(y_true, tf.float32) * weights)

def weighted_loss(y_true, y_pred, mass):
    weights = mass /K.sum(mass)
    #print(4"loss1:", K.sum( K.abs(y_pred - y_true) * weights))
    #print("loss2:", K.sum( K.abs(y_pred - y_true) * weights) / K.sum(weights))
    #return K.abs( K.sum( (y_pred - y_true) * weights) ) / K.sum(weights) 
    return K.sum( K.abs(y_pred - y_true) * weights) / K.sum(weights)

@tf.function
def curvature_trace(points, hermitian_met, kahler_pot, herm_metric, mass):
    kahler_metric = mlg.complex_math.complex_hessian(tf.math.real(kahler_pot(points)), points)
    #print("kahler_metric: ", kahler_metric)
    #inv_kahler_metric = tf.linalg.inv(kahler_metric)
    #print("inverse kahler_metric: ", inv_kahler_metric)
    curvature = mlg.complex_math.complex_hessian(tf.math.real(herm_metric(points)), points)
    #print("curvature", curvature)
    #inv_curvature = tf.linalg.inv(curvature)
    curvature_trace = tf.linalg.trace(tf.matmul(kahler_metric, curvature) )
    # print("contraction", tf.matmul(kahler_metric, curvature) )
    # print("curvature_trace: ", curvature_trace)
    return curvature_trace<jupyter_output><empty_output><jupyter_text>Generate points and load Kahler potential network<jupyter_code># Define K3 surface
z0, z1, z2, z3 = sp.symbols('z0, z1, z2, z3')
Z = [z0, z1, z2, z3]
f = z0**4 + z1**4 + z2**4 + z3**4 + 0.5*z0*z1*z2*z3
#f = z0**4 + z1**4 + z2**4 + z3**4 # Switch to Fermat surface
f
# Create training and testing set
n_pairs = 10000
HS_train = mlg.hypersurface.Hypersurface(Z, f, n_pairs)
HS_test = mlg.hypersurface.Hypersurface(Z, f, n_pairs)
train_set = mlg.tf_dataset.generate_dataset(HS_train)
test_set = mlg.tf_dataset.generate_dataset(HS_test)
train_set = train_set.shuffle(HS_train.n_points).batch(1000)
test_set = test_set.shuffle(HS_test.n_points).batch(1000)
# Load CY metric/Kahler potential
kahler_pot = tf.keras.models.load_model("K3_2_kahler_pot")
# kahler_pot = tf.keras.models.load_model("K3_fermat_kahler_pot") # Switch to fermat curve<jupyter_output>WARNING:tensorflow:No training configuration found in save file, so the model was *not* compiled. Compile it manually.<jupyter_text>O(1) This has stopped working due to no gradients occuring - $\mathcal{O}_X(2)$ is the first that works<jupyter_code># Define and train model for hermitian metric
hermitian_met = Hermitian_metric_O1(4, 1, [])
print("Hermitian metric G network defined")
c1 = tf.constant(1., shape=(1000), dtype=tf.complex64)
optimizer = tf.keras.optimizers.Adam()
loss_func = weighted_MAPE 
loss_hist = []
max_epochs = 100 
epoch = 0
#print("Beginning training")
while epoch < max_epochs:
    epoch = epoch + 1
    for step, (points, Omega_Omegabar, mass, restriction) in enumerate(train_set):
        with tf.GradientTape() as tape:
            # det_omega = nf.volume_form(points, Omega_Omegabar, mass, restriction, kahler_pot)
            curv_tr = curvature_trace(points, kahler_pot, hermitian_met, mass)
            #print("calculate curv_tr: ", curv_tr)
            #print("c1: ", c1)
            loss = loss_func(curv_tr, c1, mass)
            grads = tape.gradient(loss, hermitian_met.trainable_weights)
        optimizer.apply_gradients(zip(grads, hermitian_met.trainable_weights))
    #if epoch % 50 == 0:
    #    print("epoch %d: loss = %.7f" % (epoch, loss))
    loss_hist.append(loss)
    hermitian_met.append_loss(loss)
    print("epoch %d: loss = %.7f" % (epoch, loss))
plot_train_hist(hermitian_met, y_scale='log')<jupyter_output>Final loss:  ListWrapper([<tf.Tensor: shape=(), dtype=float32, numpy=2122.4229>])<jupyter_text>O(2) -- 2 layer network<jupyter_code># Define and train model for hermitian metric
hermitian_met = Hermitian_metric_O2(4, 1, [100])
print("Hermitian metric G network defined")
c1 = tf.constant(2., shape=(1000), dtype=tf.complex64)
optimizer = tf.keras.optimizers.Adam()
loss_func = weighted_loss  
loss_hist = []
max_epochs = 100 
epoch = 0
#print("Beginning training")
while epoch < max_epochs:
    epoch = epoch + 1
    # print("epoch: ", epoch)
    for step, (points, Omega_Omegabar, mass, restriction) in enumerate(train_set):
        # print("mass: ", mass)
        with tf.GradientTape() as tape:
            curv_tr = curvature_trace(points, kahler_pot, hermitian_met, mass)
            loss = loss_func(c1, curv_tr, mass) 
            grads = tape.gradient(loss, hermitian_met.trainable_weights)
        optimizer.apply_gradients(zip(grads, hermitian_met.trainable_weights))
    #if epoch % 50 == 0:
    #    print("epoch %d: loss = %.7f" % (epoch, loss))
    loss_hist.append(loss)
    hermitian_met.append_loss(loss)
    print("epoch %d: loss = %.7f" % (epoch, loss))
plot_train_hist(hermitian_met, y_scale='linear')<jupyter_output>Final loss:  ListWrapper([<tf.Tensor: shape=(), dtype=float32, numpy=1.9341935>])<jupyter_text>Validation/test data<jupyter_code>cal_total_loss(test_set, hermitian_met, kahler_pot, weighted_loss, c1)<jupyter_output><empty_output><jupyter_text>O(4) -- 3 layer network Seems to train reasonably well compared to the others,<jupyter_code># Define and train model for hermitian metric
hermitian_met = Hermitian_metric_O4(4, 1, [10,100])
print("Hermitian metric G network defined")
c1 = tf.constant(4., shape=(1000), dtype=tf.complex64)
optimizer = tf.keras.optimizers.Adam()
loss_func = weighted_loss 
max_epochs = 100 
epoch = 0
while epoch < max_epochs:
    epoch = epoch + 1
    for step, (points, Omega_Omegabar, mass, restriction) in enumerate(train_set):
        with tf.GradientTape() as tape:
            curv_tr = curvature_trace(points, kahler_pot, hermitian_met, mass)
            #print("calculate curv_tr: ", curv_tr)
            #print("c1: ", c1)
            loss = loss_func(c1, curv_tr, mass) 
            # print("calculated loss: ", loss)
            grads = tape.gradient(loss, hermitian_met.trainable_weights)
        optimizer.apply_gradients(zip(grads, hermitian_met.trainable_weights))
    #if epoch % 50 == 0:
    #    print("epoch %d: loss = %.7f" % (epoch, loss))
    hermitian_met.append_loss(loss)
    print("epoch %d: loss = %.7f" % (epoch, loss))
plot_train_hist(hermitian_met, y_scale='log')<jupyter_output>Final loss:  ListWrapper([<tf.Tensor: shape=(), dtype=float32, numpy=5.2278247>])<jupyter_text>Validation/test data<jupyter_code>cal_total_loss(test_set, hermitian_met, kahler_pot, weighted_loss, c1)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13054559184444026, 'fraction_numerical': 0.0205776854823485, 'mean_word_length': 3.7941176470588234, 'pattern_counts': {'":': 0, '<': 39, '<?xml version=': 0, '>': 36, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '34362', 'license': 'MIT', 'n_tokens_mistral': 3974, 'n_tokens_neox': 3837, 'n_words': 775, 'path': 'Testing K3 surface.ipynb', 'repo_name': 'yanghuihe/MLconnections', 'size': '10594'}