Draw margin curve in arcface

The effect is as follows ：

  The code is as follows ：

from math import cos, sin, pi
import numpy as np
import matplotlib.pyplot as plt

'''
# https://github.com/deepinsight/insightface/blob/master/recognition/arcface_torch/losses.py
class ArcFace(torch.nn.Module):
    """ ArcFace (https://arxiv.org/pdf/1801.07698v1.pdf):
    """
    def __init__(self, s=64.0, margin=0.5):
        super(ArcFace, self).__init__()
        self.scale = s
        self.cos_m = math.cos(margin)
        self.sin_m = math.sin(margin)
        self.theta = math.cos(math.pi - margin)
        self.sinmm = math.sin(math.pi - margin) * margin
        self.easy_margin = False


    def forward(self, logits: torch.Tensor, labels: torch.Tensor):
        index = torch.where(labels != -1)[0]
        target_logit = logits[index, labels[index].view(-1)]

        sin_theta = torch.sqrt(1.0 - torch.pow(target_logit, 2))
        cos_theta_m = target_logit * self.cos_m - sin_theta * self.sin_m  # cos(target+margin)
        if self.easy_margin:
            final_target_logit = torch.where(
                target_logit > 0, cos_theta_m, target_logit)
        else:
            final_target_logit = torch.where(
                target_logit > self.theta, cos_theta_m, target_logit - self.sinmm)

        logits[index, labels[index].view(-1)] = final_target_logit
        logits = logits * self.scale
        return logits
'''

margin = 0.5
x = np.linspace(0, pi, 100)
y1 = [(cos(i+margin) if cos(i) > 0 else cos(i)) for i in x]
y2 = [(cos(i+margin) if cos(i) > cos(pi-margin) else cos(i)-sin(pi-margin)*margin) for i in x]
l1, = plt.plot(x, y1, label='easy margin', color='red', linestyle='dashed')
l2, = plt.plot(x, y2, label='hard margin', color='green', linestyle='dotted')
plt.legend(handles=[l1,l2])
plt.show()
plt.close()

Reference resources ：

1、 Discuss easy margin

2、plot drawing

3、 Line style