Repository: sunpy/sunpy

Issue:
Input to aia_color_table should be unitful
<!-- This comments are hidden when you submit the issue so you do not need to remove them!
Please be sure to check out our contributing guidelines: https://github.com/sunpy/sunpy/blob/master/CONTRIBUTING.rst
Please be sure to check out our code of conduct:
https://github.com/sunpy/sunpy/blob/master/CODE_OF_CONDUCT.rst -->

<!-- Please have a search on our GitHub repository to see if a similar issue has already been posted.
If a similar issue is closed, have a quick look to see if you are satisfied by the resolution.
If not please go ahead and open an issue! -->

### Description
<!-- Provide a general description of the bug. -->

`aia_color_table`, the function for selecting an AIA color map, takes in a selection for the AIA wavelength. However,  the accepted input types are either a string or integer. It should be an `astropy.units.Quantity` as the channel is a wavelength and a wavelength is a physical quantity




Discussion:
Hi!
I'm working on this issue.
Thanks :-)
@wtbarnes  will you please elaborate it so that i can get this easily.Whether  unitful means unified or something else.?

Hi @mridullpandey,

To expand on the issue. 

Thee function here: https://github.com/sunpy/sunpy/blob/master/sunpy/visualization/colormaps/color_tables.py#L84
called `aia_color_table` currently takes a parameter called `wavelength` which we allow to be an integer. We want to change this to only take in and use a `astropy.units.Quantity` as we enforce this behaviour in other functions. We have a guide on this here: https://docs.sunpy.org/en/latest/dev_guide/units_quantities.html  
Hi.
So basically,what i have to do is to  change only wavelength as astropy.units.quantity form and update docstring. Is that what i have to do or some thing else ? On tables.py  its written that wavelength is based on aia_lct.pro part of SDO/AIA on SSWIDL written by Karel Schrijver (2010/04/12).  After analysing  aia_lct.pro part i found that he had taken input in angstrom , so,whether i should take unit as angstorm or meter ?.
Sorry if i missed it... and if it was much trivial....
Well astropy.units allows you to enter both units and it will convert it to the correct unit (if it can).

For the function angstrom makes the most sense. You can see an example of this here: https://github.com/sunpy/sunpy/blob/fd5b6ac409900f889260352c1ca7abeb5ae0adb3/sunpy/map/header_helper.py#L42 and the docstring for it on L81

Resolving patch:
diff --git a/sunpy/visualization/colormaps/cm.py b/sunpy/visualization/colormaps/cm.py
--- a/sunpy/visualization/colormaps/cm.py
+++ b/sunpy/visualization/colormaps/cm.py
@@ -7,32 +7,33 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
+import astropy.units as u
 from sunpy.visualization.colormaps import color_tables as ct
 
 __all__ = ['show_colormaps', 'cmlist']
 
-sdoaia94 = ct.aia_color_table(94)
-sdoaia131 = ct.aia_color_table(131)
-sdoaia171 = ct.aia_color_table(171)
-sdoaia193 = ct.aia_color_table(193)
-sdoaia211 = ct.aia_color_table(211)
-sdoaia304 = ct.aia_color_table(304)
-sdoaia335 = ct.aia_color_table(335)
-sdoaia1600 = ct.aia_color_table(1600)
-sdoaia1700 = ct.aia_color_table(1700)
-sdoaia4500 = ct.aia_color_table(4500)
-
-sohoeit171 = ct.eit_color_table(171)
-sohoeit195 = ct.eit_color_table(195)
-sohoeit284 = ct.eit_color_table(284)
-sohoeit304 = ct.eit_color_table(304)
-
-goesrsuvi94 = ct.suvi_color_table(94)
-goesrsuvi131 = ct.suvi_color_table(131)
-goesrsuvi171 = ct.suvi_color_table(171)
-goesrsuvi195 = ct.suvi_color_table(195)
-goesrsuvi284 = ct.suvi_color_table(284)
-goesrsuvi304 = ct.suvi_color_table(304)
+sdoaia94 = ct.aia_color_table(94*u.angstrom)
+sdoaia131 = ct.aia_color_table(131*u.angstrom)
+sdoaia171 = ct.aia_color_table(171*u.angstrom)
+sdoaia193 = ct.aia_color_table(193*u.angstrom)
+sdoaia211 = ct.aia_color_table(211*u.angstrom)
+sdoaia304 = ct.aia_color_table(304*u.angstrom)
+sdoaia335 = ct.aia_color_table(335*u.angstrom)
+sdoaia1600 = ct.aia_color_table(1600*u.angstrom)
+sdoaia1700 = ct.aia_color_table(1700*u.angstrom)
+sdoaia4500 = ct.aia_color_table(4500*u.angstrom)
+
+sohoeit171 = ct.eit_color_table(171*u.angstrom)
+sohoeit195 = ct.eit_color_table(195*u.angstrom)
+sohoeit284 = ct.eit_color_table(284*u.angstrom)
+sohoeit304 = ct.eit_color_table(304*u.angstrom)
+
+goesrsuvi94 = ct.suvi_color_table(94*u.angstrom)
+goesrsuvi131 = ct.suvi_color_table(131*u.angstrom)
+goesrsuvi171 = ct.suvi_color_table(171*u.angstrom)
+goesrsuvi195 = ct.suvi_color_table(195*u.angstrom)
+goesrsuvi284 = ct.suvi_color_table(284*u.angstrom)
+goesrsuvi304 = ct.suvi_color_table(304*u.angstrom)
 
 
 # The color tables below returns one of the fundamental color tables for SOHO
diff --git a/sunpy/visualization/colormaps/color_tables.py b/sunpy/visualization/colormaps/color_tables.py
--- a/sunpy/visualization/colormaps/color_tables.py
+++ b/sunpy/visualization/colormaps/color_tables.py
@@ -6,6 +6,8 @@
 import matplotlib.colors as colors
 import numpy as np
 
+import astropy.units as u
+
 __all__ = [
     'aia_color_table', 'sswidl_lasco_color_table', 'eit_color_table',
     'sxt_color_table', 'xrt_color_table', 'trace_color_table',
@@ -68,20 +70,21 @@ def paden(lst, len_, pad=0):
 c3 = ((c1 + c2 / 2.0) * 255.0 / (c1.max() + c2.max() / 2.0)).astype('f')
 
 aia_wave_dict = {
-    1600: (c3, c3, c2),
-    1700: (c1, c0, c0),
-    4500: (c0, c0, b0 / 2.0),
-    94: (c2, c3, c0),
-    131: (g0, r0, r0),
-    171: (r0, c0, b0),
-    193: (c1, c0, c2),
-    211: (c1, c0, c3),
-    304: (r0, g0, b0),
-    335: (c2, c0, c1)
+    1600*u.angstrom: (c3, c3, c2),
+    1700*u.angstrom: (c1, c0, c0),
+    4500*u.angstrom: (c0, c0, b0 / 2.0),
+    94*u.angstrom: (c2, c3, c0),
+    131*u.angstrom: (g0, r0, r0),
+    171*u.angstrom: (r0, c0, b0),
+    193*u.angstrom: (c1, c0, c2),
+    211*u.angstrom: (c1, c0, c3),
+    304*u.angstrom: (r0, g0, b0),
+    335*u.angstrom: (c2, c0, c1)
 }
 
 
-def aia_color_table(wavelength):
+@u.quantity_input
+def aia_color_table(wavelength: u.angstrom):
     """
     Returns one of the fundamental color tables for SDO AIA images.
 
@@ -90,7 +93,7 @@ def aia_color_table(wavelength):
 
     Parmeters
     ---------
-    wavelength : `int`
+    wavelength : `~astropy.units.quantity`
         Wavelength for the desired AIA color table.
     """
     try:
@@ -107,172 +110,173 @@ def aia_color_table(wavelength):
 
 
 eit_yellow_r = np.array(
-      [  0,   1,   2,   3,   5,   6,   7,   8,  10,  11,  12,  14,  15,
-        16,  17,  19,  20,  21,  22,  24,  25,  26,  28,  29,  30,  31,
-        33,  34,  35,  36,  38,  39,  40,  42,  43,  44,  45,  47,  48,
-        49,  51,  52,  53,  54,  56,  57,  58,  59,  61,  62,  63,  65,
-        66,  67,  68,  70,  71,  72,  73,  75,  76,  77,  79,  80,  81,
-        82,  84,  85,  86,  87,  89,  90,  91,  93,  94,  95,  96,  98,
-        99, 100, 102, 103, 104, 105, 107, 108, 109, 110, 112, 113, 114,
-       116, 117, 118, 119, 121, 122, 123, 124, 126, 127, 128, 130, 131,
-       132, 133, 135, 136, 137, 138, 140, 141, 142, 144, 145, 146, 147,
-       149, 150, 151, 153, 154, 155, 156, 158, 159, 160, 161, 163, 164,
-       165, 167, 168, 169, 170, 172, 173, 174, 175, 177, 178, 179, 181,
-       182, 183, 184, 186, 187, 188, 189, 191, 192, 193, 195, 196, 197,
-       198, 200, 201, 202, 204, 205, 206, 207, 209, 210, 211, 212, 214,
-       215, 216, 218, 219, 220, 221, 223, 224, 225, 226, 228, 229, 230,
-       232, 233, 234, 235, 237, 238, 239, 240, 242, 243, 244, 246, 247,
-       248, 249, 251, 252, 253, 255, 255, 255, 255, 255, 255, 255, 255,
-       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
-       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
-       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
-       255, 255, 255, 255, 255, 255, 255, 255, 255])
+    [0,   1,   2,   3,   5,   6,   7,   8,  10,  11,  12,  14,  15,
+     16,  17,  19,  20,  21,  22,  24,  25,  26,  28,  29,  30,  31,
+     33,  34,  35,  36,  38,  39,  40,  42,  43,  44,  45,  47,  48,
+     49,  51,  52,  53,  54,  56,  57,  58,  59,  61,  62,  63,  65,
+     66,  67,  68,  70,  71,  72,  73,  75,  76,  77,  79,  80,  81,
+     82,  84,  85,  86,  87,  89,  90,  91,  93,  94,  95,  96,  98,
+     99, 100, 102, 103, 104, 105, 107, 108, 109, 110, 112, 113, 114,
+     116, 117, 118, 119, 121, 122, 123, 124, 126, 127, 128, 130, 131,
+     132, 133, 135, 136, 137, 138, 140, 141, 142, 144, 145, 146, 147,
+     149, 150, 151, 153, 154, 155, 156, 158, 159, 160, 161, 163, 164,
+     165, 167, 168, 169, 170, 172, 173, 174, 175, 177, 178, 179, 181,
+     182, 183, 184, 186, 187, 188, 189, 191, 192, 193, 195, 196, 197,
+     198, 200, 201, 202, 204, 205, 206, 207, 209, 210, 211, 212, 214,
+     215, 216, 218, 219, 220, 221, 223, 224, 225, 226, 228, 229, 230,
+     232, 233, 234, 235, 237, 238, 239, 240, 242, 243, 244, 246, 247,
+     248, 249, 251, 252, 253, 255, 255, 255, 255, 255, 255, 255, 255,
+     255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+     255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+     255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+     255, 255, 255, 255, 255, 255, 255, 255, 255])
 
 eit_yellow_g = np.array(
-      [  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,
-        13,  14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,
-        26,  27,  28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,
-        39,  41,  42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,
-        53,  54,  55,  56,  57,  58,  59,  60,  61,  62,  63,  64,  65,
-        66,  67,  68,  69,  70,  71,  72,  73,  74,  75,  76,  77,  78,
-        79,  80,  82,  83,  84,  85,  86,  87,  88,  89,  90,  91,  92,
-        93,  94,  95,  96,  97,  98,  99, 100, 101, 102, 103, 104, 105,
-       106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118,
-       119, 120, 121, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132,
-       133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145,
-       146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158,
-       159, 160, 161, 162, 164, 165, 166, 167, 168, 169, 170, 171, 172,
-       173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185,
-       186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198,
-       199, 200, 201, 202, 203, 205, 206, 206, 207, 209, 209, 210, 210,
-       212, 213, 213, 215, 216, 216, 218, 219, 219, 221, 221, 222, 224,
-       224, 225, 227, 227, 228, 230, 230, 231, 231, 232, 234, 234, 235,
-       237, 237, 238, 240, 240, 241, 241, 243, 244, 244, 246, 247, 247,
-       249, 250, 250, 252, 252, 253, 255, 255, 255])
+    [0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,
+     13,  14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,
+     26,  27,  28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,
+     39,  41,  42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,
+     53,  54,  55,  56,  57,  58,  59,  60,  61,  62,  63,  64,  65,
+     66,  67,  68,  69,  70,  71,  72,  73,  74,  75,  76,  77,  78,
+     79,  80,  82,  83,  84,  85,  86,  87,  88,  89,  90,  91,  92,
+     93,  94,  95,  96,  97,  98,  99, 100, 101, 102, 103, 104, 105,
+     106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118,
+     119, 120, 121, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132,
+     133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145,
+     146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158,
+     159, 160, 161, 162, 164, 165, 166, 167, 168, 169, 170, 171, 172,
+     173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185,
+     186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198,
+     199, 200, 201, 202, 203, 205, 206, 206, 207, 209, 209, 210, 210,
+     212, 213, 213, 215, 216, 216, 218, 219, 219, 221, 221, 222, 224,
+     224, 225, 227, 227, 228, 230, 230, 231, 231, 232, 234, 234, 235,
+     237, 237, 238, 240, 240, 241, 241, 243, 244, 244, 246, 247, 247,
+     249, 250, 250, 252, 252, 253, 255, 255, 255])
 
 eit_yellow_b = np.concatenate((np.zeros(201).astype('int'),
                                np.array(
-      [  7,   7,  15,  22,  22,  30,  30,  37,  45,  45,  52,  60,  60,
-        67,  75,  75,  82,  82,  90,  97,  97, 105, 112, 112, 120, 127,
-       127, 135, 135, 142, 150, 150, 157, 165, 165, 172, 180, 180, 187,
-       187, 195, 202, 202, 210, 217, 217, 225, 232, 232, 240, 240, 247,
-       255, 255, 255])))
+    [7,   7,  15,  22,  22,  30,  30,  37,  45,  45,  52,  60,  60,
+     67,  75,  75,  82,  82,  90,  97,  97, 105, 112, 112, 120, 127,
+     127, 135, 135, 142, 150, 150, 157, 165, 165, 172, 180, 180, 187,
+     187, 195, 202, 202, 210, 217, 217, 225, 232, 232, 240, 240, 247,
+     255, 255, 255])))
 
 eit_dark_blue_r = np.concatenate((np.zeros(206).astype('int'),
                                   np.array(
-      [  9,  13,  21,  25,  25,  29,  33,  41,  49,  53,  57,  65,  69,
-        73,  77,  77,  85,  94,  98, 102, 110, 114, 118, 122, 134, 134,
-       138, 142, 146, 154, 158, 162, 166, 179, 183, 183, 187, 191, 199,
-       203, 207, 215, 223, 227, 231, 231, 235, 243, 247, 255])))
+    [9,  13,  21,  25,  25,  29,  33,  41,  49,  53,  57,  65,  69,
+     73,  77,  77,  85,  94,  98, 102, 110, 114, 118, 122, 134, 134,
+     138, 142, 146, 154, 158, 162, 166, 179, 183, 183, 187, 191, 199,
+     203, 207, 215, 223, 227, 231, 231, 235, 243, 247, 255])))
 
 eit_dark_blue_g = np.concatenate((np.zeros(128).astype('int'),
                                   np.array(
-      [  2,   2,   4,   5,   7,  12,  13,  15,  17,  20,  21,  21,  23,
-        25,  29,  31,  33,  34,  37,  39,  41,  41,  44,  47,  49,  50,
-        52,  55,  57,  60,  61,  61,  65,  66,  68,  69,  73,  76,  77,
-        79,  82,  82,  84,  85,  87,  92,  94,  95,  97, 100, 102, 103,
-       103, 105, 110, 111, 113, 114, 118, 119, 121, 122, 122, 127, 129,
-       130, 132, 135, 137, 138, 142, 145, 145, 146, 148, 150, 153, 154,
-       158, 159, 162, 164, 164, 166, 167, 170, 174, 175, 177, 180, 182,
-       183, 185, 185, 188, 191, 193, 195, 198, 199, 201, 203, 207, 207,
-       209, 211, 212, 215, 217, 219, 220, 225, 227, 227, 228, 230, 233,
-       235, 236, 239, 243, 244, 246, 246, 247, 251, 252, 255])))
+    [2,   2,   4,   5,   7,  12,  13,  15,  17,  20,  21,  21,  23,
+     25,  29,  31,  33,  34,  37,  39,  41,  41,  44,  47,  49,  50,
+     52,  55,  57,  60,  61,  61,  65,  66,  68,  69,  73,  76,  77,
+     79,  82,  82,  84,  85,  87,  92,  94,  95,  97, 100, 102, 103,
+     103, 105, 110, 111, 113, 114, 118, 119, 121, 122, 122, 127, 129,
+     130, 132, 135, 137, 138, 142, 145, 145, 146, 148, 150, 153, 154,
+     158, 159, 162, 164, 164, 166, 167, 170, 174, 175, 177, 180, 182,
+     183, 185, 185, 188, 191, 193, 195, 198, 199, 201, 203, 207, 207,
+     209, 211, 212, 215, 217, 219, 220, 225, 227, 227, 228, 230, 233,
+     235, 236, 239, 243, 244, 246, 246, 247, 251, 252, 255])))
 
 eit_dark_blue_b = np.concatenate((np.zeros(52).astype('int'),
                                   np.array(
-      [  1,   4,   5,   6,   8,   8,  10,  12,  14,  16,  18,  20,  21,
-        23,  25,  25,  28,  29,  31,  33,  35,  36,  37,  42,  43,  43,
-        44,  46,  48,  50,  51,  52,  56,  58,  59,  61,  61,  63,  65,
-        66,  69,  71,  73,  74,  75,  78,  78,  80,  81,  84,  86,  88,
-        89,  90,  93,  94,  94,  97,  99, 101, 103, 104, 105, 108, 111,
-       112, 112, 113, 116, 118, 119, 120, 124, 126, 127, 128, 131, 131,
-       132, 134, 135, 139, 141, 142, 143, 146, 147, 147, 149, 150, 154,
-       155, 157, 158, 161, 162, 164, 164, 166, 169, 170, 172, 173, 176,
-       177, 180, 181, 181, 184, 185, 187, 188, 191, 193, 195, 196, 199,
-       199, 200, 202, 203, 207, 208, 210, 211, 214, 215, 217, 217, 218,
-       222, 223, 225, 226, 229, 230, 231, 233, 233, 237, 238, 240, 241,
-       244, 245, 246, 249, 252, 252, 253, 255, 255, 255, 255, 255, 255,
-       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
-       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
-       255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
-       255, 255, 255, 255, 255, 255, 255, 255, 255])))
+    [1,   4,   5,   6,   8,   8,  10,  12,  14,  16,  18,  20,  21,
+     23,  25,  25,  28,  29,  31,  33,  35,  36,  37,  42,  43,  43,
+     44,  46,  48,  50,  51,  52,  56,  58,  59,  61,  61,  63,  65,
+     66,  69,  71,  73,  74,  75,  78,  78,  80,  81,  84,  86,  88,
+     89,  90,  93,  94,  94,  97,  99, 101, 103, 104, 105, 108, 111,
+     112, 112, 113, 116, 118, 119, 120, 124, 126, 127, 128, 131, 131,
+     132, 134, 135, 139, 141, 142, 143, 146, 147, 147, 149, 150, 154,
+     155, 157, 158, 161, 162, 164, 164, 166, 169, 170, 172, 173, 176,
+     177, 180, 181, 181, 184, 185, 187, 188, 191, 193, 195, 196, 199,
+     199, 200, 202, 203, 207, 208, 210, 211, 214, 215, 217, 217, 218,
+     222, 223, 225, 226, 229, 230, 231, 233, 233, 237, 238, 240, 241,
+     244, 245, 246, 249, 252, 252, 253, 255, 255, 255, 255, 255, 255,
+     255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+     255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+     255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
+     255, 255, 255, 255, 255, 255, 255, 255, 255])))
 
 eit_dark_green_r = np.concatenate((np.zeros(130).astype('int'),
                                    np.array(
-      [  1,   3,   4,   9,  11,  12,  14,  17,  19,  19,  20,  22,  27,
-        29,  30,  32,  35,  37,  38,  38,  41,  45,  46,  48,  50,  53,
-        54,  58,  59,  59,  62,  64,  66,  67,  71,  74,  75,  77,  80,
-        80,  82,  83,  85,  90,  91,  93,  95,  98, 100, 101, 101, 103,
-       108, 109, 111, 112, 116, 117, 119, 121, 121, 125, 127, 129, 130,
-       133, 135, 137, 140, 143, 143, 145, 146, 148, 151, 153, 156, 158,
-       161, 163, 163, 164, 166, 169, 172, 174, 175, 179, 180, 182, 183,
-       183, 187, 190, 192, 193, 196, 198, 200, 201, 206, 206, 208, 209,
-       211, 214, 216, 217, 219, 224, 225, 225, 227, 229, 232, 234, 235,
-       238, 242, 243, 245, 245, 246, 250, 251, 255])))
+    [1,   3,   4,   9,  11,  12,  14,  17,  19,  19,  20,  22,  27,
+     29,  30,  32,  35,  37,  38,  38,  41,  45,  46,  48,  50,  53,
+     54,  58,  59,  59,  62,  64,  66,  67,  71,  74,  75,  77,  80,
+     80,  82,  83,  85,  90,  91,  93,  95,  98, 100, 101, 101, 103,
+     108, 109, 111, 112, 116, 117, 119, 121, 121, 125, 127, 129, 130,
+     133, 135, 137, 140, 143, 143, 145, 146, 148, 151, 153, 156, 158,
+     161, 163, 163, 164, 166, 169, 172, 174, 175, 179, 180, 182, 183,
+     183, 187, 190, 192, 193, 196, 198, 200, 201, 206, 206, 208, 209,
+     211, 214, 216, 217, 219, 224, 225, 225, 227, 229, 232, 234, 235,
+     238, 242, 243, 245, 245, 246, 250, 251, 255])))
 
 eit_dark_green_g = np.concatenate((np.zeros(52).astype('int'),
                                    np.array(
-      [  1,   3,   4,   5,   6,   6,   8,   9,  11,  12,  14,  15,  16,
-        17,  19,  19,  21,  22,  23,  25,  26,  27,  28,  31,  32,  32,
-        33,  34,  36,  37,  38,  39,  42,  43,  44,  45,  45,  47,  48,
-        49,  51,  53,  54,  55,  56,  58,  58,  59,  60,  62,  64,  65,
-        66,  67,  69,  70,  70,  72,  73,  75,  76,  77,  78,  80,  82,
-        83,  83,  84,  86,  87,  88,  89,  92,  93,  94,  95,  97,  97,
-        98,  99, 100, 103, 104, 105, 106, 108, 109, 109, 110, 111, 114,
-       115, 116, 117, 119, 120, 121, 121, 123, 125, 126, 127, 128, 130,
-       131, 133, 134, 134, 136, 137, 138, 139, 141, 143, 144, 145, 147,
-       147, 148, 149, 150, 153, 154, 155, 156, 158, 159, 160, 160, 161,
-       164, 165, 166, 167, 169, 170, 171, 172, 172, 175, 176, 177, 178,
-       180, 181, 182, 184, 186, 186, 187, 188, 189, 191, 192, 194, 195,
-       197, 198, 198, 199, 200, 202, 204, 205, 206, 208, 209, 210, 211,
-       211, 213, 215, 216, 217, 219, 220, 221, 222, 225, 225, 226, 227,
-       228, 230, 231, 232, 233, 236, 237, 237, 238, 239, 241, 242, 243,
-       245, 247, 248, 249, 249, 250, 252, 253, 255])))
+    [1,   3,   4,   5,   6,   6,   8,   9,  11,  12,  14,  15,  16,
+     17,  19,  19,  21,  22,  23,  25,  26,  27,  28,  31,  32,  32,
+     33,  34,  36,  37,  38,  39,  42,  43,  44,  45,  45,  47,  48,
+     49,  51,  53,  54,  55,  56,  58,  58,  59,  60,  62,  64,  65,
+     66,  67,  69,  70,  70,  72,  73,  75,  76,  77,  78,  80,  82,
+     83,  83,  84,  86,  87,  88,  89,  92,  93,  94,  95,  97,  97,
+     98,  99, 100, 103, 104, 105, 106, 108, 109, 109, 110, 111, 114,
+     115, 116, 117, 119, 120, 121, 121, 123, 125, 126, 127, 128, 130,
+     131, 133, 134, 134, 136, 137, 138, 139, 141, 143, 144, 145, 147,
+     147, 148, 149, 150, 153, 154, 155, 156, 158, 159, 160, 160, 161,
+     164, 165, 166, 167, 169, 170, 171, 172, 172, 175, 176, 177, 178,
+     180, 181, 182, 184, 186, 186, 187, 188, 189, 191, 192, 194, 195,
+     197, 198, 198, 199, 200, 202, 204, 205, 206, 208, 209, 210, 211,
+     211, 213, 215, 216, 217, 219, 220, 221, 222, 225, 225, 226, 227,
+     228, 230, 231, 232, 233, 236, 237, 237, 238, 239, 241, 242, 243,
+     245, 247, 248, 249, 249, 250, 252, 253, 255])))
 
 eit_dark_green_b = np.concatenate((np.zeros(197).astype('int'), np.array(
-      [  3,  10,  17,  17,  20,  24,  27,  34,  37,  44,  48,  55,  58,
-        58,  62,  65,  72,  79,  82,  86,  93,  96,  99, 103, 103, 110,
-       117, 120, 124, 130, 134, 137, 141, 151, 151, 155, 158, 161, 168,
-       172, 175, 179, 189, 192, 192, 196, 199, 206, 210, 213, 220, 227,
-       230, 234, 234, 237, 244, 248, 255])))
+    [3,  10,  17,  17,  20,  24,  27,  34,  37,  44,  48,  55,  58,
+     58,  62,  65,  72,  79,  82,  86,  93,  96,  99, 103, 103, 110,
+     117, 120, 124, 130, 134, 137, 141, 151, 151, 155, 158, 161, 168,
+     172, 175, 179, 189, 192, 192, 196, 199, 206, 210, 213, 220, 227,
+     230, 234, 234, 237, 244, 248, 255])))
 
 eit_dark_red_r = np.concatenate((np.zeros(52).astype('int'), np.array(
-      [  1,   3,   4,   5,   6,   6,   8,   9,  11,  12,  14,  15,  16,
-        17,  19,  19,  21,  22,  23,  25,  26,  27,  28,  31,  32,  32,
-        33,  34,  36,  37,  38,  39,  42,  43,  44,  45,  45,  47,  48,
-        49,  51,  53,  54,  55,  56,  58,  58,  59,  60,  62,  64,  65,
-        66,  67,  69,  70,  70,  72,  73,  75,  76,  77,  78,  80,  82,
-        83,  83,  84,  86,  87,  88,  89,  92,  93,  94,  95,  97,  97,
-        98,  99, 100, 103, 104, 105, 106, 108, 109, 109, 110, 111, 114,
-       115, 116, 117, 119, 120, 121, 121, 123, 125, 126, 127, 128, 130,
-       131, 133, 134, 134, 136, 137, 138, 139, 141, 143, 144, 145, 147,
-       147, 148, 149, 150, 153, 154, 155, 156, 158, 159, 160, 160, 161,
-       164, 165, 166, 167, 169, 170, 171, 172, 172, 175, 176, 177, 178,
-       180, 181, 182, 184, 186, 186, 187, 188, 189, 191, 192, 194, 195,
-       197, 198, 198, 199, 200, 202, 204, 205, 206, 208, 209, 210, 211,
-       211, 213, 215, 216, 217, 219, 220, 221, 222, 225, 225, 226, 227,
-       228, 230, 231, 232, 233, 236, 237, 237, 238, 239, 241, 242, 243,
-       245, 247, 248, 249, 249, 250, 252, 253, 255])))
+    [1,   3,   4,   5,   6,   6,   8,   9,  11,  12,  14,  15,  16,
+     17,  19,  19,  21,  22,  23,  25,  26,  27,  28,  31,  32,  32,
+     33,  34,  36,  37,  38,  39,  42,  43,  44,  45,  45,  47,  48,
+     49,  51,  53,  54,  55,  56,  58,  58,  59,  60,  62,  64,  65,
+     66,  67,  69,  70,  70,  72,  73,  75,  76,  77,  78,  80,  82,
+     83,  83,  84,  86,  87,  88,  89,  92,  93,  94,  95,  97,  97,
+     98,  99, 100, 103, 104, 105, 106, 108, 109, 109, 110, 111, 114,
+     115, 116, 117, 119, 120, 121, 121, 123, 125, 126, 127, 128, 130,
+     131, 133, 134, 134, 136, 137, 138, 139, 141, 143, 144, 145, 147,
+     147, 148, 149, 150, 153, 154, 155, 156, 158, 159, 160, 160, 161,
+     164, 165, 166, 167, 169, 170, 171, 172, 172, 175, 176, 177, 178,
+     180, 181, 182, 184, 186, 186, 187, 188, 189, 191, 192, 194, 195,
+     197, 198, 198, 199, 200, 202, 204, 205, 206, 208, 209, 210, 211,
+     211, 213, 215, 216, 217, 219, 220, 221, 222, 225, 225, 226, 227,
+     228, 230, 231, 232, 233, 236, 237, 237, 238, 239, 241, 242, 243,
+     245, 247, 248, 249, 249, 250, 252, 253, 255])))
 
 eit_dark_red_g = np.concatenate((np.zeros(148).astype('int'), np.array(
-      [  1,   1,   5,   9,  11,  13,  15,  18,  20,  24,  26,  26,  30,
-        32,  34,  35,  39,  43,  45,  47,  51,  51,  52,  54,  56,  62,
-        64,  66,  68,  71,  73,  75,  75,  77,  83,  85,  86,  88,  92,
-        94,  96,  98,  98, 103, 105, 107, 109, 113, 115, 117, 120, 124,
-       124, 126, 128, 130, 134, 136, 139, 141, 145, 147, 147, 149, 151,
-       154, 158, 160, 162, 166, 168, 170, 171, 171, 175, 179, 181, 183,
-       187, 188, 190, 192, 198, 198, 200, 202, 204, 207, 209, 211, 213,
-       219, 221, 221, 222, 224, 228, 230, 232, 236, 239, 241, 243, 243,
-       245, 249, 251, 255])))
+    [1,   1,   5,   9,  11,  13,  15,  18,  20,  24,  26,  26,  30,
+     32,  34,  35,  39,  43,  45,  47,  51,  51,  52,  54,  56,  62,
+     64,  66,  68,  71,  73,  75,  75,  77,  83,  85,  86,  88,  92,
+     94,  96,  98,  98, 103, 105, 107, 109, 113, 115, 117, 120, 124,
+     124, 126, 128, 130, 134, 136, 139, 141, 145, 147, 147, 149, 151,
+     154, 158, 160, 162, 166, 168, 170, 171, 171, 175, 179, 181, 183,
+     187, 188, 190, 192, 198, 198, 200, 202, 204, 207, 209, 211, 213,
+     219, 221, 221, 222, 224, 228, 230, 232, 236, 239, 241, 243, 243,
+     245, 249, 251, 255])))
 
 eit_dark_red_b = np.concatenate((np.zeros(204).astype('int'), np.array(
-      [  3,   7,  15,  19,  27,  31,  31,  35,  39,  47,  54,  58,  62,
-        70,  74,  78,  82,  82,  90,  98, 102, 105, 113, 117, 121, 125,
-       137, 137, 141, 145, 149, 156, 160, 164, 168, 180, 184, 184, 188,
-       192, 200, 204, 207, 215, 223, 227, 231, 231, 235, 243, 247, 255])))
+    [3,   7,  15,  19,  27,  31,  31,  35,  39,  47,  54,  58,  62,
+     70,  74,  78,  82,  82,  90,  98, 102, 105, 113, 117, 121, 125,
+     137, 137, 141, 145, 149, 156, 160, 164, 168, 180, 184, 184, 188,
+     192, 200, 204, 207, 215, 223, 227, 231, 231, 235, 243, 247, 255])))
 
 
-def eit_color_table(wavelength):
+@u.quantity_input
+def eit_color_table(wavelength: u.angstrom):
     """
     Returns one of the fundamental color tables for SOHO EIT images.
     """
@@ -283,10 +287,10 @@ def eit_color_table(wavelength):
     # EIT 304 IDL Name EIT Dark Bot Red
     try:
         r, g, b = {
-            171: (eit_dark_blue_r, eit_dark_blue_g, eit_dark_blue_b),
-            195: (eit_dark_green_r, eit_dark_green_g, eit_dark_green_b),
-            284: (eit_yellow_r, eit_yellow_g, eit_yellow_b),
-            304: (eit_dark_red_r, eit_dark_red_g, eit_dark_red_b)
+            171*u.angstrom: (eit_dark_blue_r, eit_dark_blue_g, eit_dark_blue_b),
+            195*u.angstrom: (eit_dark_green_r, eit_dark_green_g, eit_dark_green_b),
+            284*u.angstrom: (eit_yellow_r, eit_yellow_g, eit_yellow_b),
+            304*u.angstrom: (eit_dark_red_r, eit_dark_red_g, eit_dark_red_b)
         }[wavelength]
     except KeyError:
         raise ValueError("Invalid EIT wavelength. Valid values are "
@@ -994,7 +998,7 @@ def sot_color_table(measurement):
             "intensity")
 
     cdict = create_cdict(r, g, b)
-    return colors.LinearSegmentedColormap( f'Hinode SOT {measurement:s}', cdict)
+    return colors.LinearSegmentedColormap(f'Hinode SOT {measurement:s}', cdict)
 
 
 def iris_sji_color_table(measurement, aialike=False):
@@ -1171,7 +1175,7 @@ def hmi_mag_color_table():
     255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255,
     255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255
 ],
-                    dtype=np.uint8)
+    dtype=np.uint8)
 
 hi1_green = np.array([
     0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
@@ -1189,7 +1193,7 @@ def hmi_mag_color_table():
     219, 220, 222, 223, 225, 227, 228, 230, 231, 233, 235, 236, 238, 239, 241,
     243, 244, 246, 247, 249, 251, 252, 254, 255
 ],
-                     dtype=np.uint8)
+    dtype=np.uint8)
 
 hi1_red = np.array([
     0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
@@ -1205,7 +1209,7 @@ def hmi_mag_color_table():
     162, 166, 170, 175, 179, 183, 187, 191, 195, 199, 203, 207, 211, 215, 219,
     223, 227, 231, 235, 239, 243, 247, 251, 255
 ],
-                   dtype=np.uint8)
+    dtype=np.uint8)
 
 hi2_blue = np.array([
     0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 4,
@@ -1224,7 +1228,7 @@ def hmi_mag_color_table():
     230, 231, 232, 233, 235, 236, 237, 238, 239, 240, 241, 242, 243, 243, 244,
     245, 246, 246, 247, 248, 249, 249, 250, 251, 252, 252, 253, 254
 ],
-                    dtype=np.uint8)
+    dtype=np.uint8)
 
 hi2_green = np.array([
     0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
@@ -1241,7 +1245,7 @@ def hmi_mag_color_table():
     155, 158, 161, 164, 167, 170, 173, 176, 180, 184, 188, 192, 197, 202, 207,
     212, 217, 222, 227, 232, 237, 242, 247, 252, 252, 253, 254
 ],
-                     dtype=np.uint8)
+    dtype=np.uint8)
 
 hi2_red = np.array([
     0, 0, 3, 6, 10, 13, 17, 20, 24, 27, 30, 33, 36, 39, 42, 45, 48, 51, 54, 57,
@@ -1262,7 +1266,7 @@ def hmi_mag_color_table():
     252, 252, 252, 252, 252, 252, 252, 252, 252, 252, 252, 252, 252, 252, 252,
     252, 252, 252, 252, 252, 253, 254
 ],
-                   dtype=np.uint8)
+    dtype=np.uint8)
 
 
 def stereo_hi_color_table(camera):
@@ -1283,19 +1287,21 @@ def create_cdict(r, g, b):
     i = np.linspace(0, 1, r0.size)
 
     cdict = {name: list(zip(i, el / 255.0, el / 255.0))
-                 for el, name in [(r, 'red'), (g, 'green'), (b, 'blue')]}
+             for el, name in [(r, 'red'), (g, 'green'), (b, 'blue')]}
     return cdict
 
 
-def suvi_color_table(wavelength):
-    """Returns one of the fundamental color tables for SUVI images.
-       SUVI uses AIA color tables.
+@u.quantity_input
+def suvi_color_table(wavelength: u.angstrom):
+    """
+    Returns one of the fundamental color tables for SUVI images.
+    SUVI uses AIA color tables.
     """
     try:
-        if wavelength == 195:
-            r, g, b = aia_wave_dict[193]
-        elif wavelength == 284:
-            r,g, b = aia_wave_dict[335]
+        if wavelength == 195*u.angstrom:
+            r, g, b = aia_wave_dict[193*u.angstrom]
+        elif wavelength == 284*u.angstrom:
+            r, g, b = aia_wave_dict[335*u.angstrom]
         else:
             r, g, b = aia_wave_dict[wavelength]
     except KeyError:


Test patch:
diff --git a/sunpy/visualization/colormaps/tests/test_cm.py b/sunpy/visualization/colormaps/tests/test_cm.py
--- a/sunpy/visualization/colormaps/tests/test_cm.py
+++ b/sunpy/visualization/colormaps/tests/test_cm.py
@@ -1,6 +1,8 @@
 import matplotlib.pyplot as plt
 import pytest
 
+import astropy.units as u
+
 import sunpy.visualization.colormaps as cm
 import sunpy.visualization.colormaps.color_tables as ct
 
@@ -33,4 +35,4 @@ def test_invalid_show_cmaps():
 def test_invalid_wavelengths(f, match):
     # Check that accessing non-existent color table values raises a ValueError
     with pytest.raises(ValueError, match=match):
-        f(-100)
+        f(-100*u.m)