unet.apl

⍝ Copy from paper: U-Net CNN in APL: Exploring Zero-Framework, Zero-Library Machine Learning (https://dl.acm.org/doi/10.1145/3589246.3595371)

W←⍬ ⋄ V←⍬ ⋄ Z←⍬ ⋄ LR←1e ̄9 ⋄ MO←0.99

FWD←{
  Z⊢←(≢W)⍴⊂⍬
  
  CV←{
    Z[⍺]←⊂⍵
    z←(,[2+⍳3]3 3⌺⍺⊃Z)+.×,[⍳3]⍺⊃W
    0⌈z⊣Z[⍺]←⊂Z[⍺],⊂z
  }
  
  CC←{
    p←2÷⍨(⍴⍺⊃Z)-⍴⍵
    ⍵,⍨(⌊p)↓(-⌈p)↓(⍺⊃Z)
  }
  
  MX←{
    ⌈⌿[2],[2 3](2 2⍴2)⌺⊃Z[⍺]←⊂⍵
  }
  
  UP←{
    Z[⍺]←⊂⍵
    s←(2× ̄1↓⍴⍵), ̄1↑⍴⍺⊃W
    s⍴0 2 1 3 4⍉⍵+.×⍺⊃W
  }
  
  C1←{
    Z[⍺]←⊂⍵
    1E ̄8+z÷[⍳2]+/z←*z-[⍳2]⌈/z←⍵+.×⍺⊃W
  }
  
  LA←{
    ⍺≥≢Z:⍵
    down←(⍺+6)∇(⍺+2)MX(⍺+1)CV(⍺+0)CV ⍵
    (⍺+2)CC(⍺+5)UP(⍺+4)CV(⍺+3)CV down
  }
  
  2 C1 1 CV 0 CV 3 LA ⍵⍴⍨3↑1,⍨⍴⍵
}

BCK←{
  Y←⍺ ⋄ Y∆←⍵
  
  ∆←{
    V[⍺]←⊂⍵+MO×(⍴⍵)⍴⍺⊃V
    W[⍺]←⊂(⍺⊃W)-LR×⍺⊃V
  }
  
  ∆CV←{
    w←,[⍳3]⊖⌽[1]0 1 3 2⍉⍺⊃W ⋄ x←⊃⍺⊃Z
    ∆z←⍵×0<1⊃⍺⊃Z
    ∆Z← ̄2⊖ ̄2⌽[1](4+2↑⍴∆z)↑∆z
    _←⍺ ∆ 3 0 1 2⍉(⍉,[⍳2]∆z)+.×,[⍳2]3 3⌺x
    w+.×⍨,[2+⍳3]3 3⌺∆Z
  }
  
  ∆CC←{
    x←⍺⊃Z ⋄ ∆z←⍵ ⋄ d←-⌊2÷⍨2↑(⍴x)-⍴∆z
    (⊃d)⊖(1⊃d)⌽[1](⍴x)↑∆z
  }
  
  ∆MX←{
    x←⍺⊃Z ⋄ ∆z←⍵
    y×x=y←(⍴x)↑2⌿2/[1]∆z
  }
  
  ∆UP←{
    w←⍺⊃W ⋄ x←⍺⊃Z ⋄ ∆z←⍵
    _←⍺ ∆(⍉,[⍳2]x)+.×,[⍳2]cz←(2 2⍴2)⌺∆z
    (,[2+⍳3]cz)+.×⍉⍪w
  }
  
  ∆C1←{
    w←⍺⊃W ⋄ x←⍺⊃Z ⋄ ∆z←⍵
    _←⍺ ∆(⍉,[⍳2]x)+.×,[⍳2]∆z
    ∆z+.×⍉w
  }
  
  ∆LA←{
    ⍺≥≢Z:⍵ ⋄ in←⍵↑[2]⍨-2÷⍨⊃⌽⍴⍵
    d←(⍺+6)∇(⍺+3)∆CV(⍺+4)∆CV(⍺+5)∆UP in
    (⍺+0)∆CV(⍺+1)∆CV(⍵ ∆CC⍨⍺+2)+(⍺+2)∆MX d
  }
  
  3 ∆LA 0 ∆CV 1 ∆CV 2 ∆C1 Y∆-(~Y),[1.5]Y
}
  
E←{
  -+⌿,⍟(⍺×⍵[;;1])+(~⍺)×⍵[;;0]
}

RUN←{
  Y←⌊0.5+nm↑⍵↓⍨2÷⍨(⍴⍵)-nm←2↑⍴Y∆←FWD ⍺
  
  Y Y∆(Y E Y∆)⊣Y BCK Y∆
}