Code

Aula 9: O fractal, pincelado, e o projeto

The programs in the class folder, highlighted, each with its file to download.

animacoes.py

Os vídeos da aula 9, feitos com o Manim Community.

671 linesDownload
"""Os vídeos da aula 9, feitos com o Manim Community.

Rode, a partir da raiz do repositório, depois de fractal.py, com o nome dos
vídeos que quer gravar, ou sem nome nenhum para gravar todos:
    python aulas/09-projeto/codigo/animacoes.py
    python aulas/09-projeto/codigo/animacoes.py zoom caixas

Os quatro vídeos, na ordem da aula:
    iteracao  repetir x ← x² + c a partir de 0: para uns c o número fica perto, para
              outros foge; na reta, a fronteira são dois pontos; no plano, é o
              conjunto de Mandelbrot; no treino, a conta repetida é o passo do gradiente
    pixel     quatro treinos da rede do projeto, um em cada par de taxas, com a perda
              passo a passo; cada um vira um pixel, e o plano se enche, de 8 por 8 a
              512 por 512 treinos
    zoom      seis janelas, uma dentro da outra, cada uma quatro vezes menor: mil vezes
              de aumento, e a fronteira continua cheia de detalhe
    caixas    a dimensão pela contagem de caixas: a linha pede 2 vezes mais caixas
              quando a caixa cai pela metade, a área 4 vezes, a fronteira 2 elevado a 1,68

Precisa do Manim Community (pip install manim), do MiKTeX e do ffmpeg, como os
vídeos das outras aulas. Cada vídeo sai em dois arquivos:
    aulas/09-projeto/videos/<nome>.mp4           1920 por 1080, 30 quadros por segundo, sem som
    aulas/09-projeto/figuras/<nome>-quadro.png   o último quadro, que o slide mostra parado

Uma opção grava só o último quadro, para conferir a composição: --quadros

As regras dos vídeos das outras aulas valem aqui. Sem título próprio, porque o
título está no slide; texto com pelo menos 30 pontos do Manim, equação com pelo
menos 38. Todo texto entra escrito e toda forma entra desenhada. Nada se
sobrepõe. E nada é sorteado: os mapas são os de fractal.py, gravados em
resultados/, e as cores são as do curso, de infra/slides.py, pela pinta() do
projeto/varredura.py.
"""
import json
import math
import shutil
import subprocess
import sys
import time
from pathlib import Path

import numpy as np
from manim import (DOWN, LEFT, RIGHT, UP, Axes, Create, DashedLine, Dot, FadeIn, FadeOut, Group, ImageMobject,
                   Line, MathTex, Rectangle, Scene, Square, TexTemplate, Text, ValueTracker, VGroup, Write, config,
                   linear, tempconfig)
from manim.constants import RESAMPLING_ALGORITHMS
from manimpango import register_font

RAIZ = Path(__file__).resolve().parents[3]
AULA = Path(__file__).resolve().parents[1]
RESULTADOS = AULA / "resultados"
sys.path.insert(0, str(Path(__file__).resolve().parent))
from figuras_slides import pinta, AZUL_LENTO, AZUL_RAPIDO, MARROM_LENTO, LARANJA_RAPIDO  # noqa: E402
from fractal import caixas, dimensao, fronteira  # noqa: E402

# As fontes do curso, as mesmas dos slides, guardadas em infra/fontes/ com as licenças
for arquivo in ("IBMPlexSans-Regular.ttf", "IBMPlexSans-SemiBold.ttf", "SourceSerif4-Semibold.ttf"):
    register_font(str(RAIZ / "infra" / "fontes" / arquivo))
SANS = "IBM Plex Sans"

# A paleta de infra/slides.py e das figuras
FUNDO, ESCURO, TEXTO, SUAVE, CINZA = "#F6F5F1", "#14213D", "#3B4658", "#5E6A7D", "#8A94A6"
AZUL, LARANJA = "#1F5FAD", "#A8500F"
CONVERGIU, DIVERGIU = AZUL, LARANJA            # as palavras "convergiu" e "divergiu", nas cores do curso

# As equações em STIX Two, a fonte das equações dos slides (infra/equacoes.py)
MODELO = TexTemplate(preamble="\\usepackage[T1]{fontenc}\n\\usepackage{amsmath}\n\\usepackage{stix2}\n")
config.tex_template = MODELO


def texto(s, tamanho=30, cor=TEXTO):
    """Texto corrido, na fonte dos slides."""
    return Text(s, font=SANS, font_size=tamanho, color=cor)


def eq(*partes, tamanho=40, cor=ESCURO):
    """Equação em LaTeX; cada parte vira um pedaço separado."""
    return MathTex(*partes, font_size=tamanho, color=cor, tex_template=MODELO)


def frase(*pedacos, tamanho=30, cor=TEXTO):
    """Uma linha que mistura texto e equação. O pedaço que começa com $ é equação, um pouco maior que o texto,
    para as letras terem a mesma altura; a pontuação depois de uma equação vai dentro dela, como em "$c,$"."""
    partes = [eq(p.strip("$"), tamanho=round(tamanho * 1.25), cor=cor) if p.startswith("$") else texto(p, tamanho, cor)
              for p in pedacos]
    return VGroup(*partes).arrange(RIGHT, buff=0.16)


def legenda(*pedacos):
    """A frase do alto do vídeo, que diz o que está acontecendo; cabe na largura com folga."""
    f = frase(*pedacos).move_to(np.array([0, 3.45, 0]))
    assert f.width < 13, pedacos
    return f


def troca(cena, atual, *pedacos):
    """Apaga a frase do alto e escreve a seguinte."""
    nova = legenda(*pedacos)
    cena.play(FadeOut(atual), run_time=0.35)
    cena.play(Write(nova), run_time=0.9)
    return nova


def br(v, casas=2):
    """Número com vírgula decimal, para texto; o menos sai como sinal de menos."""
    return f"{v:.{casas}f}".replace(".", ",").replace("-", "−")


def tex(v, casas=2):
    """Número para dentro de uma equação: vírgula decimal sem o espaço que o LaTeX põe depois dela."""
    return f"{v:.{casas}f}".replace(".", "{,}")


def mil(n):
    """Número inteiro com ponto de milhar: 262.144."""
    return f"{n:,}".replace(",", ".")


def hexa(rgb):
    """Uma cor RGB de 0 a 1 em hexadecimal, que é como o Manim recebe uma cor."""
    return "#" + "".join(f"{round(255 * v):02X}" for v in rgb)


def imagem(rgb, altura):
    """Uma matriz RGB de 0 a 1, com a linha 0 embaixo, como imagem do Manim, com os pixels nítidos."""
    img = ImageMobject((np.clip(rgb[::-1], 0, 1) * 255).astype(np.uint8))
    img.set_resampling_algorithm(RESAMPLING_ALGORITHMS["nearest"])
    img.set_height(altura)
    return img


def mascara(centro, lado, z=5):
    """Quatro retângulos da cor do fundo em volta de um quadrado: escondem o que passa dele quando a imagem cresce."""
    grande = 30
    c = np.array(centro, dtype=float)
    pecas = VGroup(
        Rectangle(width=grande, height=grande).move_to(c + np.array([0, lado / 2 + grande / 2, 0])),
        Rectangle(width=grande, height=grande).move_to(c - np.array([0, lado / 2 + grande / 2, 0])),
        Rectangle(width=grande, height=grande).move_to(c + np.array([lado / 2 + grande / 2, 0, 0])),
        Rectangle(width=grande, height=grande).move_to(c - np.array([lado / 2 + grande / 2, 0, 0])))
    pecas.set_fill(FUNDO, 1).set_stroke(width=0)
    pecas.set_z_index(z)
    return pecas


# ================================================================ 1. repetir uma conta

def iterados(c, passos):
    """Os números de x ← x² + c a partir de 0, até o passo dado ou até passar de 2.

    Com c entre −2 e 2, o número que passa de 2 nunca mais volta: x² já é maior
    que 2|x|, e cada passo o afasta mais. Por isso passar de 2 é fugir.
    """
    x, saida = 0.0, [0.0]
    for _ in range(passos):
        x = x * x + c
        saida.append(x)
        if abs(x) > 2:
            break
    return saida


def mandelbrot(n=600, iteracoes=200):
    """O conjunto de Mandelbrot em n por n pixels, com as cores dos mapas do curso.

    Cada pixel é um c = a + bi do plano, com a de −2,2 a 0,8 e b de −1,5 a 1,5.
    Repete-se z ← z² + c a partir de 0: se |z| passa de 2, a conta foge, e o
    pixel fica laranja, mais claro quanto mais cedo fugiu; se não passou em 200
    passos, fica azul.
    """
    a, b = np.meshgrid(np.linspace(-2.2, 0.8, n), np.linspace(-1.5, 1.5, n))
    c = a + 1j * b
    z = np.zeros_like(c)
    fugiu = np.full(c.shape, iteracoes)
    with np.errstate(all="ignore"):
        for k in range(iteracoes):
            z = z * z + c
            novo = (np.abs(z) > 2) & (fugiu == iteracoes)
            fugiu[novo] = k
            z[np.abs(z) > 2] = 3
    dentro = fugiu == iteracoes
    r = np.clip(1 - np.log1p(fugiu) / np.log1p(40), 0, 1)[..., None]     # 1, fugiu logo; 0, demorou
    rgb = np.where(dentro[..., None], AZUL_LENTO, MARROM_LENTO + r * (LARANJA_RAPIDO - MARROM_LENTO))
    return rgb


class Iteracao(Scene):
    """x ← x² + c a partir de 0, para quatro valores de c; depois a reta dos c, o plano, e o treino."""

    def construct(self):
        self.camera.background_color = FUNDO
        x0, x1 = -4.4, 2.2                     # a reta vai de x = −2 a x = 3; à direita, os vereditos
        cs = [(-1.0, "-1"), (0.25, "0{,}25"), (0.3, "0{,}3"), (1.0, "1")]   # o nome de c, em LaTeX
        ys = [1.9, 0.75, -0.4, -1.55]

        def px(v, y):
            return np.array([x0 + (v + 2) / 5 * (x1 - x0), y, 0])

        leg = legenda("Repetir a conta", "$x \\leftarrow x^2 + c,$", "começando em", "$x = 0$")
        guias = VGroup(*[DashedLine(px(v, ys[0] + 0.35), px(v, ys[-1] - 0.25), color=CINZA, stroke_width=1.5,
                                    dash_length=0.08) for v in range(-2, 4)])
        numeros = VGroup(*[eq(str(v), tamanho=38, cor=TEXTO).move_to(px(v, ys[-1] - 0.62))
                           for v in range(-2, 4)])
        linhas = VGroup(*[Line(px(-2, y), px(3, y), color=ESCURO, stroke_width=3) for y in ys])
        nomes = VGroup(*[eq(f"c = {n}", tamanho=40).next_to(px(-2, y), LEFT, buff=0.45) for (_, n), y in zip(cs, ys)])
        self.play(Write(leg), run_time=1.2)
        self.play(Create(linhas), Create(guias), Write(numeros), Write(nomes), run_time=1.8)

        series = [iterados(c, 14) for c, _ in cs]
        pontos = [Dot(px(0, y), radius=0.11, color=AZUL) for y in ys]
        rastros = VGroup()
        self.play(*[FadeIn(p) for p in pontos], run_time=0.5)
        passo = texto("passo 0", 30, SUAVE).move_to(np.array([5.3, 2.6, 0]))
        self.play(Write(passo), run_time=0.4)
        veredito = [None] * 4
        for k in range(1, 14):
            anims, novos = [], []
            for i, (serie, y) in enumerate(zip(series, ys)):
                if veredito[i] is not None:
                    continue
                if k >= len(serie):
                    continue
                rastros.add(Dot(pontos[i].get_center(), radius=0.05, color=AZUL, fill_opacity=0.35))
                if abs(serie[k]) <= 2:
                    anims.append(pontos[i].animate.move_to(px(serie[k], y)))
                else:                                  # passou de 2: fugiu; além de 3, sai da reta pela direita
                    veredito[i] = k
                    destino = px(serie[k], y) if serie[k] <= 3 else px(3, y) + RIGHT * 0.6
                    anims.append(pontos[i].animate.move_to(destino).set_color(LARANJA))
            novo_passo = texto(f"passo {k}", 30, SUAVE).move_to(passo)
            self.remove(passo)                       # o contador troca de uma vez, sem os dois números juntos
            self.add(rastros, novo_passo)
            passo = novo_passo
            self.play(*anims, run_time=0.55 if k < 6 else 0.4)
        # os vereditos, à direita de cada reta
        frases = []
        for i, ((c, _), y) in enumerate(zip(cs, ys)):
            if veredito[i] is not None:
                f = texto(f"foge no passo {veredito[i]}", 30, DIVERGIU)
            elif c < 0:
                f = texto("fica entre −1 e 0", 30, CONVERGIU)
            else:
                f = texto("chega perto de 0,5", 30, CONVERGIU)
            f.next_to(px(3, y), RIGHT, buff=0.95)
            assert f.get_right()[0] < 6.9, "o veredito sai da tela"
            frases.append(f)
        self.play(FadeOut(passo), *[Write(f) for f in frases], run_time=1.2)
        leg = troca(self, leg, "Fugir: passar de 2 em algum passo. Daí em diante, só cresce")
        self.wait(2.2)

        # a reta dos c: cada c pintado pela resposta
        tudo = VGroup(linhas, guias, numeros, nomes, rastros, *frases)
        self.play(FadeOut(tudo), *[FadeOut(p) for p in pontos], run_time=0.8)
        leg = troca(self, leg, "A mesma pergunta para cada", "$c$", "de", "$-2{,}5$", "a", "$1$")
        ca, cb, yc = -6.0, 5.6, 0.4

        def pc(v):
            return np.array([ca + (v + 2.5) / 3.5 * (cb - ca), yc, 0])

        eixo = Line(pc(-2.5), pc(1), color=ESCURO, stroke_width=3)
        marcas = VGroup(*[eq(n, tamanho=38, cor=TEXTO).move_to(pc(v) + DOWN * 0.75)
                          for v, n in ((-2.5, "-2{,}5"), (-2, "-2"), (-1, "-1"), (0, "0"), (0.25, "0{,}25"), (1, "1"))])
        tracos = VGroup(*[Line(pc(v) + DOWN * 0.3, pc(v) + DOWN * 0.42, color=ESCURO, stroke_width=2.5)
                          for v in (-2.5, -2, -1, 0, 0.25, 1)])
        self.play(Create(eixo), Create(tracos), Write(marcas), run_time=1.2)
        n = 140
        faixas = VGroup()
        for j in range(n):
            c = -2.5 + 3.5 * (j + 0.5) / n
            fica = all(abs(v) <= 2 for v in iterados(c, 200))
            faixas.add(Rectangle(width=(cb - ca) / n + 0.005, height=0.5, stroke_width=0)
                       .set_fill(CONVERGIU if fica else DIVERGIU, 1).move_to(pc(c)))
        self.play(FadeIn(faixas, lag_ratio=0.02), run_time=2.4)
        fronteira_ = VGroup(*[Line(pc(v) + UP * 0.45, pc(v) + DOWN * 0.3, color=ESCURO, stroke_width=5) for v in (-2, 0.25)])
        rotulos = VGroup(texto("fica", 30, CONVERGIU).move_to(pc(-0.9) + UP * 0.75),
                         texto("foge", 30, DIVERGIU).move_to(pc(-2.25) + UP * 0.75),
                         texto("foge", 30, DIVERGIU).move_to(pc(0.62) + UP * 0.75))
        self.play(Create(fronteira_), Write(rotulos), run_time=1.0)
        leg = troca(self, leg, "Na reta, a fronteira são dois pontos:", "$-2$", "e", "$0{,}25$")
        self.wait(2.4)

        # o plano: c com dois números; a mesma pergunta desenha o conjunto de Mandelbrot
        self.play(FadeOut(VGroup(eixo, marcas, tracos, faixas, fronteira_, rotulos)), run_time=0.8)
        leg = troca(self, leg, "Com", "$c$", "no plano, dois números, a fronteira é um fractal")
        mand = imagem(mandelbrot(), 5.2).move_to(np.array([-3.3, -0.35, 0]))
        nome_m = eq("z \\leftarrow z^2 + c", tamanho=44).next_to(mand, DOWN, buff=0.3)
        self.play(FadeIn(mand), Write(nome_m), run_time=1.6)
        self.wait(1.5)
        z = json.loads((RESULTADOS / "zoom.json").read_text(encoding="utf-8"))
        treino = imagem(pinta(np.load(RESULTADOS / "zoom-0.npy")), 5.2).move_to(np.array([3.3, -0.35, 0]))
        nome_t = eq("\\vec{w} \\leftarrow \\vec{w} - \\eta\\, \\nabla L(\\vec{w})", tamanho=44).next_to(treino, DOWN, buff=0.3)
        leg = troca(self, leg, "No treino, a conta repetida é o passo do gradiente")
        self.play(FadeIn(treino), Write(nome_t), run_time=1.6)
        assert z["janelas"][0] == [-3.0, 6.0, -3.0, 6.0]
        self.wait(3.0)


# ================================================================ 2. cada pixel é um treino

class Pixel(Scene):
    """Quatro treinos com a perda passo a passo; cada um vira um pixel, e o plano se enche."""

    def construct(self):
        self.camera.background_color = FUNDO
        r = json.loads((RESULTADOS / "curvas.json").read_text(encoding="utf-8"))
        S, C = 5.0, np.array([-2.7, -0.1, 0])            # o plano das taxas: lado e centro
        x0, x1 = -3.0, 6.0                                # log10 das taxas, nos dois eixos

        def no_plano(a, b):
            return C + np.array([(a - x0) / (x1 - x0) - 0.5, (b - x0) / (x1 - x0) - 0.5, 0]) * S

        quadro = Square(S, color=ESCURO, stroke_width=2.5).move_to(C)
        marcas = VGroup()
        for v in (-3, 0, 3, 6):
            n = str(v)
            marcas.add(eq(n, tamanho=38, cor=TEXTO).move_to(no_plano(v, x0) + DOWN * 0.42))
            marcas.add(eq(n, tamanho=38, cor=TEXTO).move_to(no_plano(x0, v) + LEFT * 0.42))
        nome_x = eq("\\log_{10} \\eta_0", tamanho=40).next_to(marcas, DOWN, buff=0.12)
        nome_y = eq("\\log_{10} \\eta_1", tamanho=40).rotate(math.pi / 2).next_to(quadro, LEFT, buff=0.85)

        # o painel da direita, de cima para baixo: o par de taxas (2,65), o veredito (2,05), o nome do eixo
        # vertical, o gráfico da perda, os números do passo e o nome do eixo horizontal
        eixos = Axes(x_range=[0, 100, 25], y_range=[-1, 3, 1], x_length=5.4, y_length=3.5,
                     axis_config={"color": ESCURO, "stroke_width": 2.5, "include_tip": False},
                     ).move_to(np.array([3.95, -0.75, 0]))
        rot_x = VGroup(*[eq(str(v), tamanho=38, cor=TEXTO).next_to(eixos.c2p(v, -1), DOWN, buff=0.2)
                         for v in (0, 50, 100)])
        rot_y = VGroup(*[eq(n, tamanho=38, cor=TEXTO).next_to(eixos.c2p(0, v), LEFT, buff=0.2)
                         for v, n in ((-1, "0{,}1"), (0, "1"), (1, "10"), (2, "100"), (3, "1000"))])
        nome_passo = texto("passo do treino", 30, TEXTO).next_to(rot_x, DOWN, buff=0.2)
        nome_v = frase("$L_t / L_0,$", "a perda sobre a do começo", tamanho=30).next_to(eixos, UP, buff=0.3)
        assert nome_v.get_top()[1] < 1.75 and nome_passo.get_bottom()[1] > -3.95, "o painel da direita não cabe"
        um = DashedLine(eixos.c2p(0, 0), eixos.c2p(100, 0), color=CINZA, stroke_width=2)

        leg = legenda("Cada par de taxas, uma por camada, é um treino da rede do projeto")
        self.play(Write(leg), run_time=1.2)
        self.play(Create(quadro), Write(marcas), Write(nome_x), Write(nome_y), run_time=1.6)
        self.play(Create(eixos), Write(rot_x), Write(rot_y), Write(nome_passo), Write(nome_v), Create(um), run_time=1.6)

        lado_pixel = S / 32
        pixels = VGroup()
        nomes_veredito = ["convergiu devagar", "convergiu rápido", "divergiu devagar", "divergiu rápido"]
        cores_pixel = [hexa(c) for c in (AZUL_LENTO, AZUL_RAPIDO, MARROM_LENTO, LARANJA_RAPIDO)]
        for k, ((a, b), perdas, medida) in enumerate(zip(r["pares_log10"], r["perdas"], r["medida"])):
            p = no_plano(a, b)
            alvo = Dot(p, radius=0.09, color=ESCURO)
            def taxa_tex(v):
                """A taxa para a equação: inteira com ponto de milhar a partir de 100; com uma casa abaixo disso."""
                return mil(round(v)) if v >= 100 else tex(v, 1)
            par = frase(f"$\\eta_0 = {taxa_tex(10 ** a)},$", f"$\\eta_1 = {taxa_tex(10 ** b)}$", tamanho=30)
            par.move_to(np.array([3.95, 2.65, 0]))
            self.play(FadeIn(alvo, scale=2), Write(par), run_time=0.8)
            v = [x / perdas[0] if x is not None else None for x in perdas]
            pts = []
            for t_, x in enumerate(v):
                if x is None or not math.isfinite(x) or x > 1e3:
                    pts.append(eixos.c2p(t_, 3.0))
                    break
                pts.append(eixos.c2p(t_, max(math.log10(x), -1)))
            convergiu = medida < 0
            cor = CONVERGIU if convergiu else DIVERGIU
            curva = VGroup().set_points_as_corners(pts).set_stroke(cor, 4)
            self.play(Create(curva), run_time=2.0 if len(pts) > 10 else 0.8, rate_func=linear)
            f = texto(nomes_veredito[k], 30, cor).move_to(np.array([3.95, 2.05, 0]))
            px_ = Square(lado_pixel).set_stroke(ESCURO, 1.5).set_fill(cores_pixel[k], 1).move_to(p)
            self.play(Write(f), FadeOut(alvo), FadeIn(px_), run_time=0.8)
            pixels.add(px_)
            self.wait(1.0)
            self.play(FadeOut(curva), FadeOut(f), FadeOut(par), run_time=0.5)

        # o plano se enche: 8, 16, 32, 64, 128 e 512 pixels de lado
        self.play(FadeOut(VGroup(eixos, rot_x, rot_y, nome_passo, nome_v, um)), run_time=0.6)
        leg = troca(self, leg, "Um treino por pixel, do zero, com o seu par de taxas")
        contador = VGroup()
        anterior = None
        for lado in (8, 16, 32, 64, 128, 512):
            V = np.load(RESULTADOS / (f"resolucao-{lado}.npy" if lado < 512 else "zoom-0.npy"))
            img = imagem(pinta(V), S).move_to(C)
            img.set_z_index(-1)
            novo = VGroup(texto(f"{lado} × {lado}", 40, ESCURO),
                          texto(f"{mil(lado * lado)} treinos", 30, TEXTO)).arrange(DOWN, buff=0.25)
            novo.move_to(np.array([3.6, 0.4, 0]))
            # o número velho sai antes de o novo entrar, para os dois nunca ficarem um sobre o outro
            self.play(FadeOut(contador) if anterior is not None else FadeOut(pixels), run_time=0.3)
            anims = [FadeIn(img), FadeIn(novo)]
            if anterior is not None:
                anims.append(FadeOut(anterior))
            self.play(*anims, run_time=0.8)
            anterior, contador = img, novo
            self.wait(1.2 if lado < 512 else 0.5)
        tempo = json.loads((RESULTADOS / "zoom.json").read_text(encoding="utf-8"))["segundos"][0]
        nota = texto(f"{round(tempo)} s no laptop, em 16 processos", 30, SUAVE)
        nota.next_to(contador, DOWN, buff=0.5)
        assert nota.get_right()[0] < 7.0, "a nota passa da borda"
        self.play(Write(nota), run_time=1.0)
        leg = troca(self, leg, "Azul, convergiu; laranja, divergiu. Mais claro, mais rápido")
        self.wait(3.0)


# ================================================================ 3. o zoom

class Zoom(Scene):
    """As seis janelas de fractal.py, cada uma quatro vezes menor, com o aumento e as taxas ao lado."""

    def construct(self):
        self.camera.background_color = FUNDO
        z = json.loads((RESULTADOS / "zoom.json").read_text(encoding="utf-8"))
        S, D = 6.6, np.array([-2.75, -0.4, 0])          # o lado e o centro do quadro do zoom
        imgs = [pinta(np.load(RESULTADOS / f"zoom-{k}.npy")) for k in range(len(z["janelas"]))]
        self.add(mascara(D, S))
        borda = Square(S, color=ESCURO, stroke_width=2.5).move_to(D).set_z_index(6)

        def faixa(lo, hi):
            """As duas pontas do intervalo de taxas, 10^lo e 10^hi, como texto.

            Casas decimais bastantes para as duas pontas serem diferentes quando a
            janela é estreita; abaixo de 10, pelo menos uma casa, e abaixo de 1, dois
            algarismos depois dos zeros, para 0,944 não virar 1.
            """
            a, b = 10 ** lo, 10 ** hi
            casas = max(0, int(math.ceil(-math.log10(b - a))) + 1) if b - a < 10 else 0

            def f(v):
                if v >= 1000:
                    return mil(round(v))
                if v < 1:
                    return f"{v:.2g}".replace(".", ",")      # dois algarismos significativos: 0,94; 0,001
                return br(v, max(casas, 1 if v < 10 else 0))
            return f(a), f(b)

        def painel(k):
            j = z["janelas"][k]
            e0, e1 = faixa(j[0], j[1]), faixa(j[2], j[3])
            g = VGroup(texto("aumento", 30, SUAVE), texto(f"{mil(4 ** k)} ×", 48, ESCURO),
                       frase("$\\eta_0$", f"de {e0[0]} a {e0[1]}", tamanho=30),
                       frase("$\\eta_1$", f"de {e1[0]} a {e1[1]}", tamanho=30)).arrange(DOWN, buff=0.32, aligned_edge=LEFT)
            g.move_to(np.array([4.35, -0.4, 0])).set_z_index(7)
            assert g.get_right()[0] < 7.0 and g.get_left()[0] > D[0] + S / 2 + 0.2, "o painel não cabe"
            return g

        atual = imagem(imgs[0], S).move_to(D)
        leg = legenda("A janela inteira: taxas de 0,001 a 1 milhão nas duas camadas").set_z_index(7)
        txt = painel(0)
        self.play(Write(leg), FadeIn(atual), Create(borda), Write(txt), run_time=1.6)
        self.wait(1.5)
        for k in range(1, len(imgs)):
            j0, j1 = z["janelas"][k - 1], z["janelas"][k]
            w = j0[1] - j0[0]
            F = D + np.array([((j1[0] + j1[1]) / 2 - j0[0]) / w - 0.5, ((j1[2] + j1[3]) / 2 - j0[2]) / w - 0.5, 0]) * S
            quadrado = Square(S / 4, color="black", stroke_width=4).move_to(F).set_z_index(4)
            self.play(Create(quadrado), run_time=0.8)
            proxima = imagem(imgs[k], S / 4).move_to(F)
            proxima.set_opacity(0)
            self.add(proxima)
            t = ValueTracker(0.0)

            def atualiza(_, a=atual, p=proxima, q=quadrado, F=F):
                tt = t.get_value()
                zf = 4 ** tt
                Dt = F + (D - F) * (zf - 1) / 3
                a.set_height(S * zf).move_to(Dt + zf * (D - F))
                p.set_height(S / 4 * zf).move_to(Dt)
                q.set_width(S / 4 * zf).move_to(Dt)
                # a janela nova aparece no fim, quando já cobre a maior parte do quadro: as cores de cada
                # janela se ajustam ao próprio mapa, e a borda dela ficaria visível por mais tempo
                p.set_opacity(float(np.clip((tt - 0.6) / 0.35, 0, 1)))
                q.set_stroke(opacity=float(np.clip(1 - tt / 0.5, 0, 1)))

            atual.add_updater(atualiza)
            self.play(t.animate.set_value(1.0), run_time=3.2, rate_func=lambda s: 0.5 - 0.5 * math.cos(math.pi * s))
            atual.clear_updaters()
            self.remove(atual, quadrado)
            atual = proxima
            atual.set_opacity(1)
            novo = painel(k)
            if k == 1:
                leg2 = legenda("Cada janela é quatro vezes menor que a anterior").set_z_index(7)
            elif k == len(imgs) - 1:
                leg2 = legenda("Mil vezes de aumento, e a fronteira continua cheia de detalhe").set_z_index(7)
            else:
                leg2 = None
            # o texto velho sai antes de o novo entrar
            self.play(FadeOut(txt), *([FadeOut(leg)] if leg2 is not None else []), run_time=0.3)
            self.play(FadeIn(novo), *([FadeIn(leg2)] if leg2 is not None else []), run_time=0.5)
            if leg2 is not None:
                leg = leg2
            txt = novo
            self.wait(1.4 if k < len(imgs) - 1 else 3.5)


# ================================================================ 4. a contagem de caixas

def grade_rgb(f, s, ampliar=2, tinta=(168 / 255, 80 / 255, 15 / 255), cor_fronteira=(0.08, 0.13, 0.24)):
    """A imagem da fronteira com as caixas de lado s que a tocam pintadas de laranja claro, e a grade.

    O laranja claro é o LARANJA do curso, #A8500F, misturado ao fundo.

    f: a fronteira, uma matriz de verdadeiro ou falso; a imagem sai ampliada
    `ampliar` vezes, para a grade ter linhas finas. Devolve a matriz RGB.
    """
    n = f.shape[0]
    m = n // s
    tocada = f[:m * s, :m * s].reshape(m, s, m, s).any(axis=(1, 3))
    rgb = np.ones((n, n, 3)) * np.array([0.965, 0.961, 0.945])
    caixa = np.kron(tocada, np.ones((s, s), dtype=bool))
    rgb[caixa] = 0.55 * np.array(tinta) + 0.45 * rgb[caixa]
    rgb[f] = cor_fronteira
    rgb = rgb.repeat(ampliar, 0).repeat(ampliar, 1)
    if s * ampliar >= 8:                        # as linhas da grade, só quando as caixas são grandes o bastante
        passo = s * ampliar
        rgb[::passo, :] = 0.75
        rgb[:, ::passo] = 0.75
        rgb[-1, :] = 0.75
        rgb[:, -1] = 0.75
    return rgb


class Caixas(Scene):
    """A dimensão pela contagem de caixas: a linha, a área e a fronteira da janela 3 do zoom."""

    def construct(self):
        self.camera.background_color = FUNDO
        S, C = 5.8, np.array([-3.6, -0.45, 0])
        n = 512

        # 1. a linha e a área, para calibrar
        linha_ = np.eye(n, dtype=bool)[::-1]                    # a diagonal, de baixo à esquerda ao alto à direita
        area = np.zeros((n, n), dtype=bool)
        area[128:384, 128:384] = True
        leg = legenda("Cobrir a figura com caixas e contar as que ela toca")
        self.play(Write(leg), run_time=1.0)

        def mostra(fig, tamanhos, nome, cor_nome):
            img = None
            linhas_ = VGroup()
            for i, s in enumerate(tamanhos):
                N = caixas(fig, [s])[0]
                nova = imagem(grade_rgb(fig, s), S).move_to(C)
                linha_txt = VGroup(texto(f"caixa de {s} pixels", 30, TEXTO), texto(f"{mil(N)} caixas", 30, ESCURO))
                linha_txt.arrange(RIGHT, buff=0.6)
                linha_txt.move_to(np.array([3.2, 1.6 - 0.8 * i, 0]))
                anims = [FadeIn(nova), Write(linha_txt)]
                if img is not None:
                    anims.append(FadeOut(img))
                self.play(*anims, run_time=1.0)
                img = nova
                linhas_.add(linha_txt)
                self.wait(0.6)
            return img, linhas_

        img, linhas_ = mostra(linha_, (128, 64, 32), "linha", ESCURO)
        concl = texto("caixa pela metade: 2 vezes mais caixas", 30, AZUL).move_to(np.array([3.2, -1.1, 0]))
        self.play(Write(concl), run_time=0.8)
        self.wait(1.5)
        self.play(FadeOut(img), FadeOut(linhas_), FadeOut(concl), run_time=0.6)
        img, linhas_ = mostra(area, (128, 64, 32), "área", ESCURO)
        concl = texto("caixa pela metade: 4 vezes mais caixas", 30, AZUL).move_to(np.array([3.2, -1.1, 0]))
        self.play(Write(concl), run_time=0.8)
        regra = VGroup(frase("linha: 2 =", "$2^1$", tamanho=30), frase("área: 4 =", "$2^2$", tamanho=30))
        regra.arrange(DOWN, buff=0.25, aligned_edge=LEFT).move_to(np.array([3.2, -2.4, 0]))
        leg = troca(self, leg, "O expoente é a dimensão: 1 na linha, 2 na área")
        self.play(Write(regra), run_time=1.0)
        self.wait(2.0)
        self.play(FadeOut(img), FadeOut(linhas_), FadeOut(concl), FadeOut(regra), run_time=0.6)

        # 2. a fronteira da janela 3 do zoom
        V = np.load(RESULTADOS / "zoom-3.npy")
        f = fronteira(V)
        leg = troca(self, leg, "A fronteira da janela 3 do zoom, 64 vezes ampliada")
        base = imagem(pinta(V), S).move_to(C)
        self.play(FadeIn(base), run_time=1.0)
        so_fronteira = imagem(grade_rgb(f, 1, ampliar=1), S).move_to(C)
        self.play(FadeIn(so_fronteira), run_time=1.0)
        self.remove(base)
        self.wait(1.0)

        # o i-ésimo tamanho de caixa fica em x = i + 0,5, longe do canto, para 1/16 não encostar em 2⁶
        eixos = Axes(x_range=[0, 4, 1], y_range=[6, 16, 2], x_length=4.6, y_length=3.0,
                     axis_config={"color": ESCURO, "stroke_width": 2.5, "include_tip": False},
                     x_axis_config={"include_ticks": False}).move_to(np.array([3.75, -0.95, 0]))
        tracos_x = VGroup(*[Line(eixos.c2p(i + 0.5, 6) + DOWN * 0.08, eixos.c2p(i + 0.5, 6) + UP * 0.08,
                                 color=ESCURO, stroke_width=2.5) for i in range(4)])
        eixos.add(tracos_x)
        rot_x = VGroup(*[eq(f"1/{s}", tamanho=38, cor=TEXTO).next_to(eixos.c2p(i + 0.5, 6), DOWN, buff=0.2)
                         for i, s in enumerate((16, 8, 4, 2))])
        nome_x = frase("$1/s,$", "o inverso do lado da caixa", tamanho=30).next_to(rot_x, DOWN, buff=0.15)
        assert nome_x.get_bottom()[1] > -3.9, "o nome do eixo sai da tela"
        rot_y = VGroup(*[eq(f"2^{{{v}}}", tamanho=38, cor=TEXTO).next_to(eixos.c2p(0, v), LEFT, buff=0.2)
                         for v in (6, 10, 14)])
        nome_y = frase("$N,$", "as caixas tocadas", tamanho=30).next_to(eixos, UP, buff=0.25).align_to(eixos, LEFT)
        self.play(Create(eixos), Write(rot_x), Write(nome_x), Write(rot_y), Write(nome_y), run_time=1.4)

        tamanhos = (16, 8, 4, 2)
        contagens = caixas(f, list(tamanhos))
        pontos = VGroup()
        img = so_fronteira
        contador = None
        for i, (s, N) in enumerate(zip(tamanhos, contagens)):
            nova = imagem(grade_rgb(f, s), S).move_to(C)
            ponto = Dot(eixos.c2p(i + 0.5, math.log2(N)), radius=0.09, color=LARANJA)
            novo_cont = VGroup(texto(f"caixa de {s} pixels", 30, TEXTO),
                               texto(f"{mil(N)} caixas", 30, ESCURO)).arrange(RIGHT, buff=0.6)
            novo_cont.move_to(np.array([3.6, 1.8, 0]))
            if contador is not None:
                self.play(FadeOut(contador), run_time=0.3)
            self.play(FadeIn(nova), FadeOut(img), FadeIn(ponto, scale=2), FadeIn(novo_cont), run_time=0.8)
            img, contador = nova, novo_cont
            pontos.add(ponto)
            self.wait(0.8)
        d, _ = dimensao(f, tamanhos)
        b = np.polyfit(range(4), np.log2(contagens), 1)
        reta = Line(eixos.c2p(0.2, b[1] - 0.3 * b[0]), eixos.c2p(3.8, b[1] + 3.3 * b[0]), color=AZUL, stroke_width=3)
        reta.set_z_index(-1)
        self.play(Create(reta), run_time=1.0)
        conta = VGroup(texto("caixa pela metade:", 30, TEXTO),
                       frase(f"$2^{{{tex(d)}}} \\approx {tex(2 ** d)}$", "vezes mais caixas", tamanho=30))
        conta.arrange(DOWN, buff=0.2, aligned_edge=LEFT).move_to(np.array([3.75, 2.2, 0]))
        assert conta.get_right()[0] < 7.0, "a conta passa da borda"
        assert conta.get_bottom()[1] - nome_y.get_top()[1] > 0.4, "a conta encosta no nome do eixo"
        self.play(FadeOut(contador), run_time=0.3)
        self.play(Write(conta), run_time=1.0)
        leg = troca(self, leg, f"Dimensão {br(d)}: mais que uma linha, menos que uma área")
        self.wait(3.5)


CENAS = {"iteracao": Iteracao, "pixel": Pixel, "zoom": Zoom, "caixas": Caixas}


def grava(nome, largura=1920, altura=1080, quadros=30, crf=24, so_ultimo_quadro=False):
    """Desenha a cena e grava o mp4 e o png do último quadro, em videos/ e figuras/ da aula."""
    inicio = time.perf_counter()
    pasta_videos, pasta_quadros = AULA / "videos", AULA / "figuras"
    pasta_videos.mkdir(parents=True, exist_ok=True)
    pasta_quadros.mkdir(parents=True, exist_ok=True)
    trabalho = pasta_videos / f"_manim-{nome}"
    opcoes = {"pixel_width": largura, "pixel_height": altura, "frame_rate": quadros, "media_dir": str(trabalho),
              "output_file": nome, "background_color": FUNDO, "disable_caching": True, "progress_bar": "none",
              "verbosity": "WARNING", "tex_template": MODELO}
    if so_ultimo_quadro:
        opcoes.update(save_last_frame=True, write_to_movie=False)
    with tempconfig(opcoes):
        cena = CENAS[nome]()
        cena.render()
        if so_ultimo_quadro:
            png_manim = Path(cena.renderer.file_writer.image_file_path)
        else:
            mp4 = Path(cena.renderer.file_writer.movie_file_path)
    png = pasta_quadros / f"{nome}-quadro.png"
    if so_ultimo_quadro:
        shutil.copyfile(png_manim, png)
        arquivos = [png]
    else:
        destino = pasta_videos / f"{nome}.mp4"
        subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-i", str(mp4), "-c:v", "libx264", "-crf", str(crf),
                        "-preset", "slow", "-pix_fmt", "yuv420p", "-movflags", "+faststart", "-an", str(destino)],
                       check=True)
        subprocess.run(["ffmpeg", "-y", "-loglevel", "error", "-sseof", "-0.5", "-i", str(destino),
                        "-frames:v", "1", "-update", "1", str(png)], check=True)
        arquivos = [destino, png]
    shutil.rmtree(trabalho, ignore_errors=True)
    for arquivo in arquivos:
        print(f"gravado: {arquivo} ({arquivo.stat().st_size / 1e6:.1f} MB)")
    print(f"{nome}: {time.perf_counter() - inicio:.0f} s")


if __name__ == "__main__":
    sys.stdout.reconfigure(encoding="utf-8")
    args = sys.argv[1:]
    so_quadro = "--quadros" in args
    for nome in [a for a in args if a != "--quadros"] or list(CENAS):
        grava(nome, so_ultimo_quadro=so_quadro)

figuras_slides.py

Figuras dos slides da aula 9: o fractal e o projeto.

212 linesDownload
"""Figuras dos slides da aula 9: o fractal e o projeto.

Rode, a partir da raiz do repositório, depois de fractal.py e de plano_aula7.py:
    python aulas/09-projeto/codigo/figuras_slides.py

Gera, a 200 dpi, em aulas/09-projeto/figuras/:
    gancho.png     a janela 4 do zoom em 1.024 por 1.024 treinos, sem eixos: o mapa do vídeo do
                   laço da aula 10; foi a capa desta aula até 08/11/2026
    janela.png     o mapa da janela inteira, 512 por 512 treinos, com a legenda
    janelas.png    as seis janelas do zoom, cada uma com o quadrado da seguinte
    aula7.png      a rede da aula 7 no plano das duas taxas, com a diagonal, e a perda de duas
                   taxas vizinhas por 600 passos
    dados.png      a janela inteira com o dado aleatório do artigo e com os múons do CMS

As cores, aqui, nos vídeos e no projeto/varredura.py, de onde vêm: o azul e o
laranja do curso, de infra/slides.py. Convergiu em azul, do quase branco, o
treino mais rápido, ao azul, o mais lento; divergiu em laranja, do laranja, a
explosão mais rápida, ao laranja escurecido, a mais lenta. Todo azul é mais
claro que todo laranja, e as duas regiões se separam também em escala de cinza
e para quem não distingue vermelho de verde.

O gancho leva uns 5 minutos no laptop do professor, com 16 processos, na primeira
vez; as figuras leem os .npy gravados por fractal.py e por plano_aula7.py.
"""
import json
import sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Patch, Rectangle
from matplotlib.ticker import FuncFormatter

AULA = Path(__file__).resolve().parents[1]
RAIZ = AULA.parents[1]
RESULTADOS = AULA / "resultados"
sys.path.insert(0, str(Path(__file__).resolve().parent))

plt.rcParams["mathtext.fontset"] = "stix"
ESCURO, SUAVE = "#14213d", "#5e6a7d"
VIRGULA = FuncFormatter(lambda v, _: f"{v:g}".replace(".", ",").replace("-", "−"))

# As pontas das duas escalas de cor, em RGB de 0 a 1.
# A pintura dos mapas é a do projeto, num lugar só: projeto/varredura.py, com o
# azul e o laranja do curso, de infra/slides.py.
sys.path.insert(0, str(RAIZ / "projeto"))
sys.path.insert(0, str(RAIZ / "infra"))
from slides import AZUL, LARANJA  # noqa: E402
from varredura import (CONVERGIU_LENTO as AZUL_LENTO, CONVERGIU_RAPIDO as AZUL_RAPIDO,  # noqa: E402
                       DIVERGIU_LENTO as MARROM_LENTO, DIVERGIU_RAPIDO as LARANJA_RAPIDO, pinta)

# o azul do curso clareado, para a classe "voltou" do mapa da aula 7 ficar mais
# clara que o laranja em escala de cinza
AZUL_CLARO = tuple(0.45 * np.array(AZUL_LENTO) + 0.55)


def legenda_cores(ax, ncol=2, y=-0.12, tamanho=13):
    """As quatro pontas das escalas, embaixo do mapa, fora da área dos dados."""
    itens = [Patch(color=AZUL_LENTO, label="convergiu devagar"), Patch(color=AZUL_RAPIDO, label="convergiu rápido"),
             Patch(color=MARROM_LENTO, label="divergiu devagar"), Patch(color=LARANJA_RAPIDO, label="divergiu rápido")]
    ax.legend(handles=itens, loc="upper center", bbox_to_anchor=(0.5, y), ncol=ncol, frameon=False,
              fontsize=tamanho, handlelength=1.6, columnspacing=1.6)


def eixos_taxa(ax, tamanho=16, marcas=14):
    ax.set_xlabel(r"$\log_{10}\,\eta_0$, a taxa da camada de entrada", fontsize=tamanho)
    ax.set_ylabel(r"$\log_{10}\,\eta_1$, a taxa da camada de saída", fontsize=tamanho)
    ax.xaxis.set_major_formatter(VIRGULA)
    ax.yaxis.set_major_formatter(VIRGULA)
    ax.tick_params(labelsize=marcas)


def salva(fig, nome):
    """Salva a 200 dpi, recortada em volta do que foi desenhado."""
    saida = AULA / "figuras" / nome
    saida.parent.mkdir(exist_ok=True)
    fig.savefig(saida, dpi=200, bbox_inches="tight", pad_inches=0.2)
    plt.close(fig)
    print(f"salva: {saida}")


def janela():
    """O mapa da janela inteira, a janela 0 do zoom, com a legenda embaixo."""
    z = json.loads((RESULTADOS / "zoom.json").read_text(encoding="utf-8"))
    V = np.load(RESULTADOS / "zoom-0.npy")
    fig, ax = plt.subplots(figsize=(6, 6.6))
    ax.imshow(pinta(V), origin="lower", extent=z["janelas"][0], interpolation="nearest")
    eixos_taxa(ax, 17, 15)
    legenda_cores(ax, y=-0.13, tamanho=15)
    salva(fig, "janela.png")


def janelas():
    """As seis janelas do zoom, em duas linhas de três, cada uma com o quadrado da seguinte em preto."""
    z = json.loads((RESULTADOS / "zoom.json").read_text(encoding="utf-8"))
    fig, axs = plt.subplots(2, 3, figsize=(13.5, 9.6))
    for k, ax in enumerate(axs.ravel()):
        j = z["janelas"][k]
        ax.imshow(pinta(np.load(RESULTADOS / f"zoom-{k}.npy")), origin="lower", extent=j, interpolation="nearest")
        ax.set_xticks([])
        ax.set_yticks([])
        aumento = f"{4 ** k:,}".replace(",", ".")
        ax.set_title(f"janela {k}: {aumento} ×", fontsize=21, color=ESCURO)
        if k + 1 < len(z["janelas"]):
            x0, x1, y0, y1 = z["janelas"][k + 1]
            ax.add_patch(Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, lw=2.2, ec="black"))
            ax.add_patch(Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, lw=0.9, ec="white"))
    fig.subplots_adjust(wspace=0.08, hspace=0.16)
    legenda_cores(axs[1, 1], ncol=4, y=-0.06, tamanho=19)
    salva(fig, "janelas.png")


def aula7():
    """A rede da aula 7 no plano das duas taxas, à esquerda, e a perda de três taxas vizinhas, à direita.

    À esquerda, cada pixel é um treino de 200 passos, como na aula 7, numa de três
    classes no passo 200 (plano_aula7.py): voltou, a perda abaixo da de antes do
    treino; em trânsito, ainda acima dela; explodiu. A diagonal tracejada é a taxa
    única da aula 7; o quadrado preto, a janela da faixa.
    À direita, duas taxas vizinhas da janela da faixa, na mesma linha, a segunda
    0,9% maior, treinadas por 600 passos: a perda das duas salta para dezenas e
    fica oscilando, e no passo 200 uma está abaixo da perda de antes e a outra
    acima. É a causa do granulado.
    Cores do curso: azul-claro, laranja e laranja escurecido, três níveis de
    cinza bem separados; nas curvas, cor, traço e marcador diferentes.
    """
    r = json.loads((RESULTADOS / "plano-aula7.json").read_text(encoding="utf-8"))
    antes = r["perda_antes"]
    cores3 = np.array([AZUL_CLARO, LARANJA_RAPIDO, MARROM_LENTO])
    fig, axs = plt.subplots(1, 2, figsize=(13, 6.1))
    for ax in axs:
        ax.set_box_aspect(0.75)

    # à esquerda, o plano no passo 200
    ax = axs[0]
    c = np.load(RESULTADOS / "plano-aula7-plano.npy")
    j = r["plano"]["janela"]
    ax.imshow(cores3[c], origin="lower", extent=j, interpolation="nearest", aspect="auto")
    eixos_taxa(ax, 17, 15)
    ax.plot([j[0], j[1]], [j[0], j[1]], ls="--", lw=2, color="black")
    x0, x1, y0, y1 = r["faixa"]["janela"]
    ax.add_patch(Rectangle((x0, y0), x1 - x0, y1 - y0, fill=False, lw=2.2, ec="black"))
    ax.set_title("o plano, no passo 200", fontsize=18, color=ESCURO)
    itens = [Patch(color=cores3[0], label="voltou"), Patch(color=cores3[1], label="em trânsito"),
             Patch(color=cores3[2], label="explodiu"),
             plt.Line2D([], [], ls="--", lw=2, color="black", label=r"$\eta_0 = \eta_1$, a aula 7")]
    ax.legend(handles=itens, loc="upper center", bbox_to_anchor=(0.5, -0.2), ncol=2, frameon=False,
              fontsize=16, handlelength=2.2, columnspacing=1.4)

    # à direita, duas taxas vizinhas, 600 passos
    ax = axs[1]
    passos = np.arange(len(r["vizinhas"][0]["perdas"]))
    for vz, (cor, traco, marca) in zip(r["vizinhas"][4:6], [(LARANJA, "-", "s"), (AZUL, "-", "o")]):
        k = vz["voltou_no_passo"]
        quando = (f"abaixo de 2,32 no passo {k}" if k else f"acima de 2,32 até o passo {passos[-1]}")
        ax.semilogy(passos, vz["perdas"], color=cor, ls=traco, lw=1.6, marker=marca, ms=8, markevery=50,
                    label=rf"$\eta_0$ = {vz['eta0']:.2f}: {quando}".replace(".", ",", 1))
    ax.axhline(antes, color=ESCURO, lw=2, ls=(0, (1, 1.5)), label="a perda de antes do treino, 2,32")
    ax.axvline(200, color=SUAVE, lw=2, ls=(0, (6, 3)), label="o passo 200, o da aula 7")
    ax.set_xlim(0, passos[-1])
    ax.set_ylim(0.8, 120)
    ax.set_xlabel("passo do treino", fontsize=17)
    ax.set_ylabel("perda", fontsize=17)
    ax.yaxis.set_major_formatter(VIRGULA)
    ax.tick_params(labelsize=15)
    ax.set_title("duas taxas vizinhas, 600 passos", fontsize=18, color=ESCURO)
    ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.2), ncol=1, frameon=False, fontsize=16,
              handlelength=2.6)
    fig.subplots_adjust(wspace=0.3, bottom=0.3, top=0.93, left=0.08, right=0.97)
    salva(fig, "aula7.png")


def dados():
    """A janela inteira com o dado aleatório do artigo e com os múons, lado a lado (projeto/figuras/)."""
    fig, axs = plt.subplots(1, 2, figsize=(13, 7.2))
    for ax, nome, titulo in zip(axs, ("aleatorio", "muons"), ("dado aleatório, como no artigo", "72 eventos do CMS")):
        reg = json.loads((RAIZ / "projeto" / "figuras" / f"janela-{nome}.json").read_text(encoding="utf-8"))
        V = np.load(RAIZ / "projeto" / "figuras" / f"janela-{nome}.npy")
        ax.imshow(pinta(V), origin="lower", extent=reg["janela"], interpolation="nearest")
        eixos_taxa(ax, 17, 15)
        ax.set_title(titulo, fontsize=19, color=ESCURO)
    fig.subplots_adjust(wspace=0.28, bottom=0.2)
    itens = [Patch(color=AZUL_LENTO, label="convergiu devagar"), Patch(color=AZUL_RAPIDO, label="convergiu rápido"),
             Patch(color=MARROM_LENTO, label="divergiu devagar"), Patch(color=LARANJA_RAPIDO, label="divergiu rápido")]
    fig.legend(handles=itens, loc="lower center", ncol=4, frameon=False, fontsize=17, bbox_to_anchor=(0.5, 0.02))
    salva(fig, "dados.png")


def gancho():
    """A janela 4 do zoom em 1.024 por 1.024 treinos, sem eixos: o mapa do laço da aula 10."""
    from fractal import janela_do_zoom, mapa, rede_e_dados
    destino = RESULTADOS / "gancho.npy"
    if not destino.exists():
        np.save(destino, mapa(janela_do_zoom(4), 1024, rede_e_dados()))
    V = np.load(destino)
    fig = plt.figure(figsize=(1024 / 200, 1024 / 200), dpi=200)
    ax = fig.add_axes([0, 0, 1, 1])
    ax.imshow(pinta(V), origin="lower", interpolation="nearest")
    ax.axis("off")
    saida = AULA / "figuras" / "gancho.png"
    fig.savefig(saida, dpi=200)
    plt.close(fig)
    print(f"salva: {saida}")


if __name__ == "__main__":
    sys.stdout.reconfigure(encoding="utf-8")
    pedidas = sys.argv[1:] or ["janela", "janelas", "aula7", "dados", "gancho"]
    for nome in pedidas:
        globals()[nome]()

fractal.py

Aula 9: o mapa das duas taxas, o zoom na fronteira e a dimensão dela.

328 linesDownload
"""Aula 9: o mapa das duas taxas, o zoom na fronteira e a dimensão dela.

Rode, a partir da raiz do repositório, depois de infra/baixar_dados.py:
    python aulas/09-projeto/codigo/fractal.py curvas
    python aulas/09-projeto/codigo/fractal.py resolucoes
    python aulas/09-projeto/codigo/fractal.py zoom
    python aulas/09-projeto/codigo/fractal.py dimensao
    python aulas/09-projeto/codigo/fractal.py esquerda
    python aulas/09-projeto/codigo/fractal.py precisao

A rede, os dados e o treino são os do projeto, em projeto/varredura.py, e este
programa só os chama: 8 entradas, a energia e o momento dos dois múons; 8
neurônios com a tanh; uma saída, o logaritmo da massa; sem viés; 72 eventos do
CMS, tantos quanto os pesos da rede; gradiente descendente com todos os
exemplos, 100 passos, uma taxa para cada camada; semente 0. Cada pixel de um
mapa é um treino inteiro, do zero, com um par de taxas.

O que cada parte grava em resultados/:
- curvas: a perda de quatro treinos, passo a passo, para o vídeo em que cada
  pixel é um treino (curvas.json);
- resolucoes: o mapa da janela inteira com 8, 16, 32, 64 e 128 pixels de lado,
  para o mesmo vídeo (resolucoes.json e resolucao-N.npy);
- zoom: seis janelas, uma dentro da outra, cada uma quatro vezes menor, com
  512 por 512 treinos cada, para o vídeo do zoom (zoom.json e zoom-K.npy);
- dimensao: a dimensão da fronteira em cada janela do zoom, pela contagem de
  caixas, e a mediana, como no artigo (dimensao.json).
- esquerda: a fronteira de cima, à esquerda do mapa: o limiar da regressão
  linear da aula 2, 2 / λ, e os treinos que explodem e voltam logo acima dele
  (esquerda.json);
- precisao: a janela do zoom do projeto em float32 e em float64, comparadas
  pixel a pixel (precisao.json).

Os .npy, a imagem em números, ficam fora do git: se refazem pelo .json ao lado.

As varreduras se repartem entre os núcleos do processador: cada processo treina
um pedaço dos pixels, com a função treinar do projeto. No laptop do professor,
com 16 núcleos, a janela de 512 por 512 leva uns 73 segundos, e tudo, uns 8 minutos.
"""
import os

# Cada processo usa um núcleo só na álgebra linear: são 16 processos, um por
# núcleo, e mais linhas de execução por processo só disputariam os núcleos.
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
os.environ.setdefault("OMP_NUM_THREADS", "1")

import json  # noqa: E402
import sys  # noqa: E402
import time  # noqa: E402
from concurrent.futures import ProcessPoolExecutor  # noqa: E402
from pathlib import Path  # noqa: E402

import numpy as np  # noqa: E402

RAIZ = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(RAIZ / "projeto"))
from varredura import carregar_muons, medida, sha256, treinar, ARQUIVO_MUONS  # noqa: E402

RESULTADOS = Path(__file__).resolve().parents[1] / "resultados"
SEMENTE, LARGURA, PASSOS = 0, 8, 100
JANELA = (-3.0, 6.0, -3.0, 6.0)          # a janela inteira do projeto, em log10 da taxa
PROCESSOS = os.cpu_count() or 4

# O zoom: o centro de cada janela, em log10 das duas taxas, e a largura da
# primeira. Cada janela seguinte tem um quarto da largura da anterior e fica
# dentro dela: a última é 4⁵ = 1.024 vezes menor que a primeira. Os centros
# foram escolhidos olhando cada mapa, sobre a fronteira de cima, onde há
# filamentos e redemoinhos: é escolha à mão, como a do artigo.
CENTROS = [(1.5, 1.5), (1.1, 1.65), (1.0, 1.72), (1.12, 1.665), (1.085, 1.668), (1.0825, 1.6715)]
LARGURA_0 = 9.0


def rede_e_dados():
    """Os pesos iniciais e os 72 eventos do projeto, como em varredura.py com a semente 0.

    Devolve X (72 por 8), y (72), W0 (8 por 8) e W1 (8). A ordem dos sorteios é
    a mesma de varredura.main(): primeiro os pesos, depois os eventos.
    """
    rng = np.random.default_rng(SEMENTE)
    W0, W1 = rng.standard_normal((8, LARGURA)), rng.standard_normal(LARGURA)
    X, y = carregar_muons(8 * LARGURA + LARGURA, SEMENTE)
    return X, y, W0, W1


def _pedaco(args):
    """Treina um pedaço dos pixels num processo e devolve a medida de cada treino."""
    X, y, W0, W1, taxa0, taxa1, tipo = args
    return medida(treinar(X, y, W0, W1, taxa0, taxa1, PASSOS, tipo))


def mapa(janela, resolucao, dados, tipo=np.float64):
    """O mapa de uma janela: resolucao por resolucao treinos, um por par de taxas.

    janela: (x0, x1, y0, y1) em log10 da taxa; x é a camada de entrada, y a de saída.
    Devolve a matriz de medidas, linha = taxa de saída, coluna = taxa de entrada:
    negativa onde o treino convergiu, positiva onde divergiu (medida() do projeto).
    tipo: a precisão dos números, float64, como no artigo, ou float32, para comparar.
    """
    tx = np.logspace(janela[0], janela[1], resolucao)
    ty = np.logspace(janela[2], janela[3], resolucao)
    taxa0, taxa1 = (m.ravel() for m in np.meshgrid(tx, ty))
    bloco = -(-taxa0.size // (4 * PROCESSOS))                 # quatro pedaços por processo
    tarefas = [(*dados, taxa0[i:i + bloco], taxa1[i:i + bloco], tipo) for i in range(0, taxa0.size, bloco)]
    with ProcessPoolExecutor(PROCESSOS) as ex:
        partes = list(ex.map(_pedaco, tarefas))
    return np.concatenate(partes).reshape(resolucao, resolucao)


def janela_do_zoom(k):
    """A janela k do zoom, (x0, x1, y0, y1), quadrada em log10 das taxas."""
    cx, cy = CENTROS[k]
    m = LARGURA_0 / 4 ** k / 2
    return (cx - m, cx + m, cy - m, cy + m)


def fronteira(V):
    """Os pixels da fronteira: os que têm um vizinho, de lado ou de cima, do outro lado.

    V é o mapa de medidas; o lado é o sinal, convergiu ou divergiu. Devolve uma
    matriz de verdadeiro ou falso do tamanho do mapa.
    """
    c = V < 0
    f = np.zeros_like(c)
    f[:, 1:] |= c[:, 1:] != c[:, :-1]
    f[:, :-1] |= c[:, 1:] != c[:, :-1]
    f[1:, :] |= c[1:, :] != c[:-1, :]
    f[:-1, :] |= c[1:, :] != c[:-1, :]
    return f


def caixas(f, tamanhos):
    """Quantas caixas de cada tamanho, em pixels, tocam a fronteira f.

    A imagem se corta numa grade de caixas de lado s; conta-se a caixa que tem
    pelo menos um pixel da fronteira. Uma linha precisa de 2 vezes mais caixas
    quando s cai pela metade; uma área, de 4 vezes; a fronteira, de 2 elevado à
    dimensão dela.
    """
    n = f.shape[0]
    contagem = []
    for s in tamanhos:
        m = n // s
        g = f[:m * s, :m * s].reshape(m, s, m, s).any(axis=(1, 3))
        contagem.append(int(g.sum()))
    return contagem


def dimensao(f, tamanhos=(2, 4, 8, 16)):
    """A dimensão pela contagem de caixas: a inclinação de log N contra log(1/s).

    Devolve a inclinação e as contagens. A reta se ajusta por mínimos quadrados.
    O tamanho 1, o pixel, fica de fora: nele, a fronteira desenhada tem a
    espessura de um pixel, e a conta mede o desenho, não a fronteira. As
    caixas de 32 e 64 pixels também: nas janelas de dentro do zoom, quase
    todas tocam a fronteira, a contagem para de crescer e a reta entorta.
    """
    N = caixas(f, tamanhos)
    inclinacao = np.polyfit(np.log(1 / np.array(tamanhos)), np.log(N), 1)[0]
    return float(inclinacao), N


def grava(nome, dados):
    RESULTADOS.mkdir(exist_ok=True)
    (RESULTADOS / nome).write_text(json.dumps(dados, indent=1, ensure_ascii=False), encoding="utf-8")
    print(f"gravado: {RESULTADOS / nome}")


def origem():
    return {"semente": SEMENTE, "largura": LARGURA, "passos": PASSOS, "exemplos": 8 * LARGURA + LARGURA,
            "dados": "dados/dimuon.csv", "sha256": sha256(ARQUIVO_MUONS)}


def faz_curvas(dados):
    """A perda, passo a passo, de quatro treinos da janela inteira: dois que convergem e dois que divergem.

    Os pares, em log10 das taxas (entrada; saída), foram escolhidos no mapa: um
    no meio da região que converge, um perto da fronteira, um logo acima dela e
    um longe, no alto.
    """
    pares = [(0.0, 0.0), (1.0, 1.5), (1.0, 2.2), (1.0, 4.5)]
    X, y, W0, W1 = dados
    t0 = np.array([10 ** a for a, _ in pares])
    t1 = np.array([10 ** b for _, b in pares])
    perdas = treinar(X, y, W0, W1, t0, t1, PASSOS)
    v = medida(perdas)
    grava("curvas.json", {**origem(), "pares_log10": pares, "perdas": np.where(np.isfinite(perdas), perdas, None).tolist(),
                          "medida": v.tolist()})


def faz_resolucoes(dados):
    """O mapa da janela inteira com 8, 16, 32, 64 e 128 pixels de lado: o vídeo enche o plano aos poucos."""
    tempos = {}
    for r in (8, 16, 32, 64, 128):
        inicio = time.perf_counter()
        V = mapa(JANELA, r, dados)
        tempos[r] = round(time.perf_counter() - inicio, 1)
        np.save(RESULTADOS / f"resolucao-{r}.npy", V)
        print(f"{r} por {r}: {tempos[r]} s, convergiram {np.mean(V < 0) * 100:.0f}%")
    grava("resolucoes.json", {**origem(), "janela": JANELA, "segundos": tempos, "processos": PROCESSOS})


def faz_zoom(dados, resolucao=512):
    """As seis janelas do zoom, com o tempo de cada uma."""
    janelas, tempos, fracoes = [], [], []
    for k in range(len(CENTROS)):
        j = janela_do_zoom(k)
        inicio = time.perf_counter()
        V = mapa(j, resolucao, dados)
        tempos.append(round(time.perf_counter() - inicio, 1))
        np.save(RESULTADOS / f"zoom-{k}.npy", V)
        janelas.append(j)
        fracoes.append(round(float(np.mean(V < 0)), 4))
        print(f"janela {k}: largura {j[1] - j[0]:.6f}, {tempos[-1]} s, convergiram {fracoes[-1] * 100:.0f}%")
    grava("zoom.json", {**origem(), "resolucao": resolucao, "centros": CENTROS, "largura_0": LARGURA_0,
                        "janelas": janelas, "segundos": tempos, "fracao_convergiu": fracoes, "processos": PROCESSOS})


def faz_dimensao():
    """A dimensão da fronteira em cada janela do zoom e a mediana delas, como no artigo."""
    r = json.loads((RESULTADOS / "zoom.json").read_text(encoding="utf-8"))
    tamanhos = (2, 4, 8, 16, 32, 64)        # as seis contagens vão para o JSON; a reta usa as quatro primeiras
    por_janela = []
    for k in range(len(r["janelas"])):
        f = fronteira(np.load(RESULTADOS / f"zoom-{k}.npy"))
        d, _ = dimensao(f)
        N = caixas(f, tamanhos)
        por_janela.append({"dimensao": round(d, 4), "caixas": N})
        print(f"janela {k}: dimensão {d:.2f}, caixas {N}")
    mediana = float(np.median([p["dimensao"] for p in por_janela]))
    print(f"mediana: {mediana:.2f}")
    grava("dimensao.json", {"tamanhos": tamanhos, "tamanhos_da_reta": (2, 4, 8, 16), "janelas": por_janela,
                            "mediana": round(mediana, 4)})


def faz_esquerda(dados, eta0=1e-3, etas1=(50.0, 52.0, 60.0, 70.0, 80.0)):
    """A fronteira de cima, à esquerda do mapa, onde a camada de entrada quase não anda.

    Com η₀ muito pequena, só a camada de saída treina, e o treino é a regressão
    linear da aula 2 sobre as saídas fixas da camada oculta, h: a perda é uma
    tigela, e o passo diverge acima de η₁ = 2 / λ, com λ o maior autovalor da
    hessiana, (2/n) hᵀh / 8², a conta do limiar 0,513 da reta da aula 2.
    O mapa mostra mais que a conta: logo acima do limiar, uma linha fina de
    treinos que divergem; acima dela, uma faixa em que a perda explode e volta.
    Na explosão, o gradiente fica tão grande que a camada de entrada se mexe,
    mesmo com η₀ = 0,001, o maior autovalor cai abaixo de 2 / η₁, e o treino
    converge; mais acima, a explosão vence. Os treinos com η₀ = 0,001 e as
    taxas de etas1 mostram isso com a função de leitura, perda_e_gradientes.
    """
    from varredura import perda_e_gradientes
    X, y, W0, W1 = dados
    n, d = X.shape

    def maior_autovalor(W):
        h = np.tanh(np.sqrt(2 / d) * X @ W)
        return float(np.linalg.eigvalsh(2 / n * h.T @ h / W1.size ** 2).max())

    lam = maior_autovalor(W0)
    treinos = []
    for eta1 in etas1:
        A, B, perdas = W0.copy(), W1.copy(), []
        with np.errstate(all="ignore"):
            for _ in range(PASSOS):
                L, g0, g1 = perda_e_gradientes(X, y, A, B)
                perdas.append(float(L))
                if not np.isfinite(L) or L > 1e6:
                    break
                A -= eta0 * g0
                B -= eta1 * g1
        divergiu = not np.isfinite(perdas[-1]) or perdas[-1] > 1e6 or perdas[-1] > perdas[0]
        treinos.append({"eta1": eta1, "pico": round(max(perdas), 1), "final": round(perdas[-1], 4),
                        "divergiu": bool(divergiu), "lambda_final": None if divergiu else round(maior_autovalor(A), 5),
                        "W0_mudou": None if divergiu else round(float(np.linalg.norm(A - W0) / np.linalg.norm(W0)), 3)})
        print(treinos[-1])
    # no mapa da janela inteira: a classe de cada linha nas colunas com η₀ abaixo de 10⁻²
    linha = {}
    if (RESULTADOS / "zoom-0.npy").exists():
        V = np.load(RESULTADOS / "zoom-0.npy")
        ty = np.linspace(JANELA[2], JANELA[3], V.shape[0])
        colunas = np.linspace(JANELA[0], JANELA[1], V.shape[1]) < -2
        fracao = (V[:, colunas] > 0).mean(1)            # em cada linha, a fração dos treinos que divergem
        fina = [i for i in range(len(ty)) if ty[i] < 1.8 and fracao[i] > 0.5 and fracao[i - 1] == 0 and fracao[i + 1] == 0]
        topo = float(ty[np.where(fracao == 0)[0].max()])
        linha = {"linha_fina_log10": [round(float(ty[i]), 3) for i in fina],
                 "linha_fina_fracao": [round(float(fracao[i]), 2) for i in fina],
                 "ultimo_que_converge_log10": round(topo, 3)}
    grava("esquerda.json", {**origem(), "eta0": eta0, "maior_autovalor": round(lam, 5),
                            "limiar": round(2 / lam, 2), "treinos": treinos, "no_mapa": linha})


def faz_precisao(dados, janela=(0.6, 1.6, 1.45, 1.95), resolucao=256):
    """A mesma janela em float32 e em float64: quantos pixels caem na mesma classe.

    A janela é a do zoom-muons de projeto/figuras/. Devolve, em precisao.json, a
    fração de pixels com a mesma classe, convergiu ou divergiu, nas duas
    precisões, e a fração dos que trocam que estão em cima da fronteira do mapa
    em float64. Também o espaçamento relativo entre taxas vizinhas da janela e o
    menor espaçamento que cada precisão distingue.
    """
    V32 = mapa(janela, resolucao, dados, np.float32)
    V64 = mapa(janela, resolucao, dados, np.float64)
    troca = (V32 < 0) != (V64 < 0)
    grava("precisao.json", {**origem(), "janela": janela, "resolucao": resolucao,
                            "mesma_classe": round(float(np.mean(~troca)), 4),
                            "trocas_na_fronteira": round(float(np.mean(fronteira(V64)[troca])), 3),
                            "espacamento_relativo": float(np.log(10) * (janela[1] - janela[0]) / (resolucao - 1)),
                            "float32_distingue": float(np.finfo(np.float32).eps),
                            "float64_distingue": float(np.finfo(np.float64).eps)})


if __name__ == "__main__":
    sys.stdout.reconfigure(encoding="utf-8")
    RESULTADOS.mkdir(exist_ok=True)
    o_que = sys.argv[1] if len(sys.argv) > 1 else "tudo"
    if o_que == "dimensao":
        faz_dimensao()
        sys.exit()
    dados = rede_e_dados()
    if o_que in ("curvas", "tudo"):
        faz_curvas(dados)
    if o_que in ("resolucoes", "tudo"):
        faz_resolucoes(dados)
    if o_que in ("zoom", "tudo"):
        faz_zoom(dados)
    if o_que == "tudo":
        faz_dimensao()
    if o_que in ("esquerda", "tudo"):
        faz_esquerda(dados)
    if o_que in ("precisao", "tudo"):
        faz_precisao(dados)

plano_aula7.py

Aula 9: a rede pequena da aula 7 no plano das duas taxas, e por que a faixa dela é granulada.

235 linesDownload
"""Aula 9: a rede pequena da aula 7 no plano das duas taxas, e por que a faixa dela é granulada.

Rode, a partir da raiz do repositório:
    python aulas/09-projeto/codigo/plano_aula7.py

A aula 7 treinou a rede de 784 entradas, 16 neurônios com a tanh e 10 saídas
nos 256 primeiros exemplos do MNIST, 200 passos sem lotes, com uma taxa só, e
achou, entre 3,86 e 7,55, uma faixa em que taxas vizinhas dão respostas
diferentes. Aqui cada camada ganha a sua taxa: η₀ para a camada oculta, pesos e
vieses, e η₁ para a de saída. A taxa única da aula 7 é a diagonal do plano,
η₀ = η₁.

Cada pixel é um treino, e a perda de cada passo fica guardada. No passo k, o
treino está numa de três classes:
- voltou: a perda está abaixo da de antes do treino, 2,3220;
- em trânsito: a perda está acima dela, mas nunca passou de 1.000;
- explodiu: a perda passou de 1.000, ou virou infinita, em algum passo até k.
"Voltou" é o "melhorou a rede" da aula 7, e a regra do artigo, a média dos 20
últimos passos abaixo da perda do começo, dá a mesma classe em mais de 99% dos
pixels.

O programa grava em resultados/plano-aula7.json:
1. o plano, η₀ e η₁ de 10⁻² a 10^2,5, 128 por 128 treinos de 200 passos, com
   as três classes, e a classe de cada pixel em plano-aula7-plano.npy;
2. a janela da faixa, η₀ de 10^0,6 a 10^1,1 e η₁ de 10^0,55 a 10^1,05, 128 por
   128 treinos de 600 passos: quantos já voltaram nos passos 200, 400 e 600, e
   quantos pares de vizinhos estão em classes diferentes;
3. a perda passo a passo de seis taxas vizinhas da janela, na mesma linha: o
   retrato do trânsito;
4. a diagonal da aula 7, as mesmas 601 taxas de 2 a 10 da figura do zoom de lá,
   agora com 600 passos: quantas pioram a rede e quantas vezes a resposta troca
   entre taxas vizinhas, nos passos 200, 400 e 600.
Só a parte 4, que leva segundos, sem refazer as outras:
    python aulas/09-projeto/codigo/plano_aula7.py diagonal

O resultado explica a faixa granulada. Nos primeiros passos, a perda salta da
inicial, 2,32, para dezenas, entre 37 e 62 nas seis vizinhas, e depois volta
devagar, oscilando, e cada taxa volta num passo diferente, sem ordem entre
vizinhas. O mapa num passo fixo pega umas
já de volta e outras ainda no caminho: com mais passos, mais voltam, e o mapa
continua granulado.

Para treinar milhares de redes de uma vez, a conta se reorganiza sem mudar o
resultado. O passo do gradiente descendente na camada oculta soma aos pesos uma
combinação das linhas de X, as imagens; então X W₀ᵀ + b₀, a entrada da tanh,
muda por K D, com K = X Xᵀ + 1, uma matriz de 256 por 256 calculada uma vez, e
D, a derivada da perda em relação a essa entrada. Assim a rede nunca multiplica
pelos 784 pixels depois do primeiro passo. A conferência no início compara três
taxas da diagonal com a função erro_final da aula 7.

No laptop do professor, o plano leva cerca de 2 minutos e meio, e a janela de
600 passos, uns 8 minutos; os tempos medidos ficam no .json.
"""
import json
import math
import sys
import time
from pathlib import Path

import numpy as np
import torch

RAIZ = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(RAIZ / "aulas" / "07-pytorch" / "codigo"))
sys.path.insert(0, str(RAIZ / "aulas" / "06-mnist-numpy" / "codigo"))
from mnist import carrega, versao  # noqa: E402
from taxa import N_EXEMPLOS, PASSOS, SEMENTE, erro_final, perda, pesos_iniciais  # noqa: E402

RESULTADOS = Path(__file__).resolve().parents[1] / "resultados"
PLANO = (-2.0, 2.5, -2.0, 2.5)
FAIXA = (0.6, 1.1, 0.55, 1.05)
RESOLUCAO = 128
PASSOS_FAIXA = 600
LINHA, COLUNAS = 64, range(40, 46)        # os seis vizinhos da janela da faixa: linha 64, colunas 40 a 45
VOLTOU, TRANSITO, EXPLODIU = 0, 1, 2


def treina_muitas(taxa0, taxa1, X, Y, P, K, passos=PASSOS):
    """Treina G cópias da rede pequena, uma para cada par de taxas, e devolve a perda de cada passo.

    taxa0, taxa1: tensores de G taxas, da camada oculta e da de saída;
    X, Y: as 256 imagens e os rótulos em one-hot, 256 por 10;
    P: os pesos iniciais da aula 7; K: X Xᵀ + 1.
    Devolve uma matriz G por passos + 1: a perda antes do primeiro passo e
    depois de cada um. A conta é a de taxa.py da aula 7, reorganizada.
    """
    G, n = taxa0.numel(), X.shape[0]
    Z = (X @ P[0].T + P[1])[None].expand(G, n, -1).clone()    # a entrada da tanh, G por 256 por 16
    W1 = P[2].T[None].expand(G, -1, -1).clone()                # os pesos de saída, G por 16 por 10
    b1 = P[3][None, None].expand(G, 1, -1).clone()             # os vieses de saída, G por 1 por 10
    e0, e1 = taxa0.view(G, 1, 1), taxa1.view(G, 1, 1)
    perdas = torch.empty(G, passos + 1, dtype=torch.float64)
    for k in range(passos + 1):
        h = torch.tanh(Z)
        ls = torch.log_softmax(torch.baddbmm(b1, h, W1), 2)    # o logaritmo da softmax
        perdas[:, k] = -(ls * Y).sum(2).mean(1)                # a entropia cruzada de cada rede
        if k == passos:
            break
        d = (ls.exp() - Y) / n                                 # a derivada da perda em relação à saída
        grad_W1 = torch.bmm(h.transpose(1, 2), d)
        grad_b1 = d.sum(1, keepdim=True)
        D = torch.bmm(d, W1.transpose(1, 2)) * (1 - h * h)     # a regra da cadeia até a entrada da tanh
        Z -= e0 * torch.matmul(K, D)                           # o passo da camada oculta, visto em Z
        W1 -= e1 * grad_W1
        b1 -= e1 * grad_b1
    return perdas.numpy()


def classes(perdas, k, antes):
    """A classe de cada treino no passo k: voltou, em trânsito ou explodiu (ver o começo do arquivo)."""
    ate_k = perdas[:, :k + 1]
    explodiu = (~np.isfinite(ate_k) | (ate_k > 1e3)).any(1)
    voltou = ~explodiu & (perdas[:, k] < antes)
    return np.where(explodiu, EXPLODIU, np.where(voltou, VOLTOU, TRANSITO))


def taxas(janela, resolucao=RESOLUCAO):
    """Os pares de taxas de uma janela em log10, linha por linha: linha = η₁, coluna = η₀."""
    tx = torch.logspace(janela[0], janela[1], resolucao, dtype=torch.float64)
    ty = torch.logspace(janela[2], janela[3], resolucao, dtype=torch.float64)
    T0, T1 = torch.meshgrid(tx, ty, indexing="xy")
    return T0.reshape(-1), T1.reshape(-1)


def varre(janela, passos, X, Y, P, K, bloco=2048):
    """A perda de cada passo de todos os treinos da janela, uma linha por pixel."""
    t0, t1 = taxas(janela)
    return np.concatenate([treina_muitas(t0[i:i + bloco], t1[i:i + bloco], X, Y, P, K, passos)
                           for i in range(0, t0.numel(), bloco)])


def vizinhos_diferentes(c):
    """A fração dos pares de pixels vizinhos, lado a lado, que estão em classes diferentes."""
    m = c.reshape(RESOLUCAO, RESOLUCAO)
    return float(np.mean(m[:, 1:] != m[:, :-1]))


def diagonal(X, Y, P, K, antes):
    """A parte 4: as 601 taxas de 2 a 10 da aula 7, com η₀ = η₁, por 600 passos.

    Devolve, para os passos 200, 400 e 600, quantas das 601 pioram a rede, a
    perda acima da de antes, e quantas vezes a resposta, melhora ou piora, troca
    entre taxas vizinhas. No passo 200, a aula 7 conta 383 e 77 trocas; aqui a
    ordem das somas é outra, e na faixa o treino amplifica a diferença do
    arredondamento, então as contagens podem diferir um pouco.
    """
    eta = torch.linspace(2, 10, 601, dtype=torch.float64)
    H = treina_muitas(eta, eta, X, Y, P, K, PASSOS_FAIXA)
    saida = {"taxas": [2, 10, 601], "passos": PASSOS_FAIXA}
    for k in (200, 400, 600):
        piora = classes(H, k, antes) != VOLTOU
        saida[str(k)] = {"pioram": int(piora.sum()), "trocas": int(np.sum(piora[1:] != piora[:-1]))}
    print("diagonal:", saida)
    return saida


if __name__ == "__main__":
    sys.stdout.reconfigure(encoding="utf-8")
    torch.set_num_threads(16)
    X, t = carrega()["treino"]
    X, t = torch.tensor(X[:N_EXEMPLOS]), torch.tensor(t[:N_EXEMPLOS])
    Y = torch.nn.functional.one_hot(t, 10).double()
    P = pesos_iniciais()
    K = X @ X.T + 1.0
    with torch.no_grad():
        antes = perda(P, X, t).item()
    if sys.argv[1:] == ["diagonal"]:
        resultado = json.loads((RESULTADOS / "plano-aula7.json").read_text(encoding="utf-8"))
        resultado["diagonal"] = diagonal(X, Y, P, K, antes)
        (RESULTADOS / "plano-aula7.json").write_text(json.dumps(resultado, ensure_ascii=False), encoding="utf-8")
        sys.exit()

    # Conferência: na diagonal, a conta reorganizada tem de dar o resultado da aula 7.
    # Só com taxas de antes da faixa: dentro dela, a ordem diferente das somas muda
    # o arredondamento na 16ª casa, o treino amplifica a diferença, e os dois números
    # finais se separam. A taxa 5 fica registrada para mostrar isso.
    def final(eta):
        h = treina_muitas(torch.tensor([eta], dtype=torch.float64), torch.tensor([eta], dtype=torch.float64), X, Y, P, K)
        return math.inf if classes(h, PASSOS, antes)[0] == EXPLODIU else float(h[0, -1])
    conferencia = []
    for eta in (0.5, 1.0, 2.0):
        a, b = erro_final(eta, X, t, P), final(eta)
        conferencia.append({"taxa": eta, "aula7": a, "aqui": b})
        print(f"taxa {eta}: aula 7 {a:.10f}, aqui {b:.10f}")
        assert abs(a - b) < 1e-8, "a conta reorganizada não confere com a da aula 7"
    caos = {"taxa": 5.0, "aula7": erro_final(5.0, X, t, P), "aqui": final(5.0)}
    print(f"taxa 5, na faixa: aula 7 {caos['aula7']:.4f}, aqui {caos['aqui']:.4f}")

    RESULTADOS.mkdir(exist_ok=True)
    resultado = {"versao_dos_dados": versao(), "semente": SEMENTE, "exemplos": N_EXEMPLOS, "perda_antes": antes,
                 "resolucao": RESOLUCAO, "conferencia": conferencia, "caos": caos}

    # 1. o plano, 200 passos, como na aula 7
    inicio = time.perf_counter()
    H = varre(PLANO, PASSOS, X, Y, P, K)
    c = classes(H, PASSOS, antes)
    np.save(RESULTADOS / "plano-aula7-plano.npy", c.reshape(RESOLUCAO, RESOLUCAO).astype(np.int8))
    v = H / antes                                               # a regra do artigo, para comparar
    artigo = (np.nan_to_num(v[:, -20:], nan=np.inf, posinf=np.inf).mean(1) < 1) & (c != EXPLODIU)
    resultado["plano"] = {"janela": PLANO, "passos": PASSOS, "segundos": round(time.perf_counter() - inicio, 1),
                          "fracao": {nome: round(float(np.mean(c == k)), 4)
                                     for nome, k in (("voltou", VOLTOU), ("transito", TRANSITO),
                                                     ("explodiu", EXPLODIU))},
                          "regra_do_artigo_igual": round(float(np.mean(artigo == (c == VOLTOU))), 4)}
    print("plano:", resultado["plano"])

    # 2. a janela da faixa, 600 passos
    inicio = time.perf_counter()
    H = varre(FAIXA, PASSOS_FAIXA, X, Y, P, K)
    por_passo = {}
    for k in (200, 400, 600):
        ck = classes(H, k, antes)
        por_passo[str(k)] = {"voltou": round(float(np.mean(ck == VOLTOU)), 4),
                             "vizinhos_diferentes": round(vizinhos_diferentes(ck), 4)}
        np.save(RESULTADOS / f"plano-aula7-faixa-{k}.npy", ck.reshape(RESOLUCAO, RESOLUCAO).astype(np.int8))
    c200, c600 = classes(H, 200, antes), classes(H, 600, antes)
    resultado["faixa"] = {"janela": FAIXA, "passos": PASSOS_FAIXA, "segundos": round(time.perf_counter() - inicio, 1),
                          "por_passo": por_passo,
                          "voltou_depois_do_200": int(np.sum((c200 != VOLTOU) & (c600 == VOLTOU))),
                          "voltou_e_subiu_de_novo": int(np.sum((c200 == VOLTOU) & (c600 != VOLTOU)))}
    print("faixa:", resultado["faixa"])

    # 3. seis vizinhas na mesma linha da janela da faixa
    t0, t1 = taxas(FAIXA)
    idx = [LINHA * RESOLUCAO + j for j in COLUNAS]
    resultado["vizinhas"] = [{"eta0": float(t0[i]), "eta1": float(t1[i]),
                              "perdas": [round(float(x), 5) for x in H[i]],
                              "voltou_no_passo": next((k for k in range(1, PASSOS_FAIXA + 1) if H[i, k] < antes
                                                       and H[i, :k].max() > antes), None)}
                             for i in idx]
    resultado["diagonal"] = diagonal(X, Y, P, K, antes)
    print("vizinhas: voltam nos passos", [vz["voltou_no_passo"] for vz in resultado["vizinhas"]],
          "| pico", [round(max(vz["perdas"]), 1) for vz in resultado["vizinhas"]])
    (RESULTADOS / "plano-aula7.json").write_text(json.dumps(resultado, ensure_ascii=False), encoding="utf-8")
    print(f"gravado: {RESULTADOS / 'plano-aula7.json'}")