元素周期表及元素性质

from manim import *
import elementy.elements as elements

frame_width = config["frame_width"]
frame_height = config["frame_height"]

Elements = elements.get_elements()

def narrator(a):
    for i in a:
        globals()[f't{str(i)}']=Text(a[i], font = "STZhongsong").to_edge(DOWN).scale(0.6)

class Element(VGroup):
    def __init__(
        self,
        **kwargs
    ):
        super().__init__(**kwargs)
        square = RoundedRectangle(height=2, width=2, corner_radius=0.3, fill_opacity=0.8, fill_color=GRAY)
        self.add(square)
        self.square = square

    def add_content(self, current_content = None, next_content = None, is_number = True):

        if is_number:
            font_size = 15
        else:
            font_size = 25

        if current_content is not None:
            mob_current_content = Text(current_content, font_size=font_size).move_to(self)
            self.add(mob_current_content)
            self.current_content = mob_current_content

        if next_content is not None:
            mob_next_content = Text(next_content, font_size=font_size).move_to(self).flip(axis=[1,0,0]).set_opacity(0)
            self.add(mob_next_content)
            self.next_content = mob_next_content
        
    @property
    def flip(self):
        anim =  AnimationGroup(self.square.animate.flip(axis=[1,0,0]).set_color(self.target_color),
                              self.current_content.animate.flip(axis=[1,0,0]).set_opacity(0),
                              self.next_content.animate.flip(axis=[1,0,0]).set_opacity(1))
        return anim

    def get_target_color(self, value, min, max, color1, color2):
        target_value = (value-min)/(max-min)
        color = interpolate_color(color1, color2, target_value)
        self.target_color = color

class ElementsPeriodicTable(VGroup):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        row = 5
        col = 18
        elements = VGroup()
        for i in range(row):
            for j in range(col):
                elements.add(Element())
        
        elements.arrange_in_grid(row, col, buff=0.5)
        elements.set_width(frame_width*0.95)
        del elements.submobjects[1:17]
        del elements.submobjects[4:14]
        del elements.submobjects[12:22]

        self.elements = elements
        self.add(elements)

        self.origin()

    def origin(self):
        n = 0
        symbol = []
        for element in self.elements:
            element.add_content(current_content=Elements[n].symbol, is_number=False)
            symbol.append(Elements[n].symbol)
            n += 1
        self.symbol = symbol

    def to_property(self, property_list, color1, color2):
        max_value = max(property_list)
        min_value = min(property_list)

        self.get_full_graph()
        n = 0
        for element in self.elements:
            value = property_list[n]
            element.add_content(next_content=str(round(value, 2)))
            element.get_target_color(value, min_value, max_value, color1, color2)
            n += 1
        anim1 = LaggedStart(*[ele.flip for ele in self.elements], run_time=3)
        anim2 = Wait(3)
        anim = Succession(anim1, anim2)

        return anim

    @property
    def to_affinity(self):
        property_list = [Elements[i].electron_affinity for i in range(len(self.elements))]
        print(property_list)
        self.property_list = property_list
        self.graph_color = PURPLE
        return self.to_property(property_list, PURPLE_A, PURPLE_E)
        
    
    @property
    def to_ionization(self):
        property_list = [Elements[i].ionisation_energies[0] for i in range(len(self.elements))]
        self.property_list = property_list
        self.graph_color = BLUE
        return self.to_property(property_list, BLUE_A, BLUE_E)

    @property
    def to_radius(self):
        property_list = [Elements[i].radius_covalent for i in range(len(self.elements))]
        self.property_list = property_list
        self.graph_color = GREEN
        return self.to_property(property_list, GREEN_A, GREEN_E)
    
    def get_graph(self):
        property_list = self.property_list
        period = self.period
        symbol = self.symbol
        if period == 1:
            ele_index = [0, 2]
        elif period == 2:
            ele_index = [2, 10]
        elif period == 3:
            ele_index = [10, 18]
        elif period == 4:
            ele_index = [18, 36]
        elif period == 5:
            ele_index = [36, 54]
        period_property_list = property_list[ele_index[0]:ele_index[1]]
        period_symbol = symbol[ele_index[0]:ele_index[1]]
        max_value = max(period_property_list)
        min_value = min(period_property_list)
        plane = NumberPlane(
            x_range = (0, len(period_property_list)),
            y_range = (min_value, max_value),
            x_length = frame_width*0.8,
            y_length = frame_height*0.6,
            axis_config={"include_numbers": True}
        )
        graph = plane.plot_line_graph(
            x_values = [i + 1 for i in range(len(self.elements))],
            y_values = period_property_list,
            line_color = self.graph_color
        )
        graph.center()
        n = 0
        graph_vg = VGroup(graph)
        for dot in graph["vertex_dots"]:
            text = Text(period_symbol[n], font_size=30).next_to(dot, UP)
            graph_vg.add(text)
            n += 1
        self.graph = graph_vg
    
    def set_period(self, period):
        self.period = period
        self.get_graph()

    def get_full_graph(self):
        property_list = self.property_list
        symbol = self.symbol
        max_value = max(property_list)
        min_value = min(property_list)
        plane = NumberPlane(
            x_range = (0, len(property_list)),
            y_range = (min_value, max_value),
            x_length = frame_width*0.8,
            y_length = frame_height*0.6,
            axis_config={"include_numbers": True}
        )
        graph = plane.plot_line_graph(
            x_values = [i + 1 for i in range(len(self.elements))],
            y_values = property_list,
            line_color = self.graph_color
        )
        graph.center()
        # n = 0
        # graph_vg = VGroup(graph)
        # for dot in graph["vertex_dots"]:
        #     text = Text(symbol[n], font_size=15).next_to(dot, UP)
        #     graph_vg.add(text)
        #     n += 1
        self.full_graph = graph
        return graph

class AtomicOrbitalTemplate(TexTemplate):
    def __init__(self,**kwargs):
        super().__init__(**kwargs)
        self.add_to_preamble("\\usepackage{tikzorbital}")

class AtomicOrbital(VGroup):
    def __init__(
        self,
        orbital_name: str,
        orbital_name_tex: str,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.template = AtomicOrbitalTemplate()
        orbital = MathTex("\\begin{tikzpicture} \\orbital{%s}  \\end{tikzpicture}" % orbital_name, tex_template = self.template, stroke_width=2, fill_color=BLUE, fill_opacity=0.6, **kwargs)
        name = MathTex(orbital_name_tex)
        name.next_to(orbital, DOWN)
        self.orbital = orbital
        self.name = name
        self.add(self.orbital)
        self.add(self.name)

class AtomicElectronsArrangement(VGroup):
    def __init__(self, orbital_name, orbital_label, electrons_num, **kwargs):
        super().__init__(**kwargs)
        if orbital_name == 's':
            self.orbital_lines_buff = 2
            self.orbital_num = 1
            self.orbital_name_shift_distance = 0
            orbital = AtomicOrbital('s', orbital_label)
        elif orbital_name == 'p':
            self.orbital_lines_buff = 6
            self.orbital_num = 3
            self.orbital_name_shift_distance = 0.3
            orbital = AtomicOrbital('py', orbital_label)
        elif orbital_name == 'd':
            self.orbital_lines_buff = 6
            self.orbital_num = 5
            self.orbital_name_shift_distance = 0.8
            orbital = AtomicOrbital('dyz', orbital_label)
        lines = VGroup(*[Line(LEFT/2, RIGHT/2) for i in range(self.orbital_num)])
        lines.arrange(buff=0.5)
        orbital.next_to(lines[0], self.orbital_lines_buff*LEFT)
        orbital.name.shift(self.orbital_name_shift_distance*DOWN)
        self.orbital = orbital
        self.add(lines, orbital)
        self.center()
        dots = VGroup()
        if electrons_num > self.orbital_num:
            for i in range(self.orbital_num):
                dot = Dot(color=BLUE).move_to(lines[i]).shift(0.2*UP + 0.3*LEFT)
                arrow = MathTex(r'\uparrow').set_color(BLUE).scale(0.8).next_to(dot, 0.5*UP)
                dots.add(VGroup(dot ,arrow))
            for i in range(electrons_num - self.orbital_num):
                dot = Dot(color=RED).move_to(lines[i]).shift(0.2*UP + 0.3*RIGHT)
                arrow = MathTex(r'\downarrow').set_color(RED).scale(0.8).next_to(dot, 0.5*UP)
                dots.add(VGroup(dot ,arrow))
        else:
            for i in range(electrons_num):
                dot = Dot(color=BLUE).move_to(lines[i]).shift(0.2*UP + 0.3*LEFT)
                arrow = MathTex(r'\uparrow').set_color(BLUE).scale(0.8).next_to(dot, 0.5*UP)
                dots.add(VGroup(dot ,arrow))
        dots.set_opacity(0)
        self.add(dots)
        self.dots = dots

        self.show_dots = self.anim_show_dots()

    def anim_show_dots(self):
        anim = Succession(*[ApplyMethod(dot.set_opacity, 1) for dot in self.dots])  # bug：使用 animate 会使物件 opacity 变为 1 的同时回到屏幕中心
        return anim
    
class AnimIonization(VGroup):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        self.Be_B = self.mob_Be_B()
        self.N_O = self.mob_N_O()
    
    def mob_Be_B(self):
        text = VGroup(Text('Be', font_size=30), Text('B', font_size=30)).arrange(buff=5)
        text.shift(2*DOWN)
        Be_e = AtomicElectronsArrangement('s', '2s', 2).scale(0.8).next_to(text[0], 6*UP)
        B_e1 = AtomicElectronsArrangement('s', '2s', 2).scale(0.8)
        B_e2 = AtomicElectronsArrangement('p', '2p', 1).scale(0.8)
        B_e = VGroup(B_e1, B_e2).arrange(direction=UP, buff=1).next_to(text[1], 2*UP)
        self.show_Be_B_dots = Succession(Be_e.show_dots, B_e1.show_dots, B_e2.show_dots)
        mob_vg = VGroup(text, Be_e, B_e)
        return mob_vg
    
    def mob_N_O(self):
        text = VGroup(Text('N', font_size=30), Text('O', font_size=30)).arrange(buff=5)
        text.shift(DOWN)
        N_e = AtomicElectronsArrangement('p', '2p', 3).scale(0.6).next_to(text[0], 3*UP)
        O_e = AtomicElectronsArrangement('p', '2p', 4).scale(0.6).next_to(text[1], 3*UP)
        self.show_N_O_dots = Succession(N_e.show_dots, O_e.show_dots)
        mob_vg = VGroup(text, N_e, O_e)
        return mob_vg

class AnimAtomRadius(VGroup):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        self.two_atoms = self.mob_two_atoms()
        self.mob_atom()
        self.shrink = self.anim_shrink()
        self.grow = self.anim_grow()

    def anim_atom_wave(self, mob, radius=2):
        wave = Circle(radius=radius).set_color(color=[BLUE_B, BLUE_D])
        anim = Broadcast(wave, focal_point=mob.get_center(), n_mobs=4, initial_opacity=0.8, run_time=4)
        return anim

    def mob_two_atoms(self):
        atom1 = Circle(radius=2, fill_opacity=0.8).set_color(color=[BLUE_A, BLUE_E])
        atom2 = Circle(radius=2, fill_opacity=0.8).set_color(color=[BLUE_E, BLUE_A]).shift(3*RIGHT)
        line = DashedLine(0.75*LEFT, 0.75*RIGHT).shift(0.75*RIGHT)
        vertical_line = Line(3*UP, 3*DOWN, color=BLUE)
        atoms = VGroup(atom1, atom2).center()
        self.show_two_atoms = AnimationGroup(self.anim_atom_wave(atom1), self.anim_atom_wave(atom2))
        lines = VGroup(line, vertical_line)
        self.show_lines_of_two_atom = Succession(DrawBorderThenFill(lines), FadeOut(lines))
        return atoms
    
    def mob_atom(self):
        single_atom = Circle(radius=2, fill_opacity=0.8).set_color(color=[BLUE_A, BLUE_E]).set_z_index(-2)
        nucleus = Circle(radius=0.3, fill_opacity=0.8).set_color(color=[GREEN_A, GREEN_E]).set_z_index(2)
        atom = VGroup(single_atom, nucleus)
        single_atom.save_state()
        self.single_atom = single_atom
        self.nucleus = nucleus
        self.atom = atom
    
    def anim_shrink(self):
        dots_lines = VGroup()
        for i in range(4):
            dot = Dot().set_color(color=[RED, PURPLE])
            distance = 0.3 + np.random.random()
            angle = np.random.random()*2*PI
            dot.shift([distance*np.cos(angle), distance*np.sin(angle), 0])
            line = DashedLine(dot.get_center(), self.atom.get_center(), dashed_ratio=0.5).set_opacity(0.6)
            rotate_angle = angle_of_vector(dot.get_center())
            arrow = MathTex(r'\leftarrow').scale(0.8).shift(line.get_length()/2*RIGHT).rotate(rotate_angle, about_point=ORIGIN)
            dot_line = VGroup(dot, line, arrow).add_updater(lambda z: z.rotate(0.04, about_point=ORIGIN))
            dots_lines.add(dot_line)
            self.shrink_dots_lines = dots_lines


        anim = Succession(*[Succession(AnimationGroup(FadeIn(dot_line), ApplyMethod(self.single_atom.scale, 0.9)), Wait(2)) for dot_line in dots_lines])
        return anim
    
    def anim_grow(self):
        dots = VGroup()
        dots_lines = VGroup()
        for i in range(4):
            dot = Dot().set_color(color=[RED, PURPLE])
            distance = 0.5 + np.random.random()*1.5
            angle = np.random.random()*2*PI
            dot.shift([distance*np.cos(angle), distance*np.sin(angle), 0])
            dot.add_updater(lambda z: z.rotate(0.04, about_point=ORIGIN))
            lines = VGroup()
            for pre_dot in dots:
                def generate_updater(line, dot1, dot2):
                    def updater(line):
                        line.become(DashedLine(dot1.get_center(), dot2.get_center(), dashed_ratio=0.5).set_opacity(0.6))
                    return updater
                single_line = DashedLine(dot.get_center(), pre_dot.get_center(), dashed_ratio=0.5).set_opacity(0.6)
                single_line.add_updater(generate_updater(single_line, dot, pre_dot))
                lines.add(single_line)
            dot_line = VGroup(dot, lines)
            dots.add(dot)
            dots_lines.add(dot_line)
            self.grow_dots_lines = dots_lines


        anim = Succession(Restore(self.single_atom), *[Succession(AnimationGroup(FadeIn(dot_line, suspend_mobject_updating=False, run_time=0.01), ApplyMethod(self.single_atom.scale, 1.1)), Wait(2)) for dot_line in dots_lines])
        return anim